vkfetch_rs/
lib.rs

1pub mod ascii_art;
2pub mod device;
3pub mod vendor;
4
5use ascii_art::{BRIGHT_GREEN, BRIGHT_RED, BRIGHT_YELLOW};
6use ash::{self, Entry, Instance, vk};
7use device::Device;
8use std::{
9    error::Error,
10    ffi::CStr,
11    io::{self, Write},
12};
13use vt::enable_virtual_terminal_processing;
14
15const BOLD: &str = "\x1B[1m";
16const RESET: &str = "\x1B[0m";
17const DIM: &str = "\x1B[90m";
18const WRAP_OFF: &str = "\x1B[?7l";
19const WRAP_ON: &str = "\x1B[?7h";
20const ALIGNMENT: &str = "    ";
21const EMPTY: &str = "";
22
23/// Fetches and prints information for a given physical device.
24pub fn fetch_device(
25    instance: &Instance,
26    device_handle: vk::PhysicalDevice,
27) -> Result<(), Box<dyn Error>> {
28    let _ = enable_virtual_terminal_processing();
29    let use_ansi = is_ansi_supported();
30
31    let device = Device::new(instance, device_handle);
32    let vendor = device.vendor;
33    let art = vendor.get_ascii_art_with_ansi(use_ansi);
34    let blank_art = " ".repeat(vendor.ascii_art_width());
35
36    let accent = if use_ansi {
37        vendor.get_alternative_style()[0]
38    } else {
39        EMPTY
40    };
41    let info = get_device_info(&device, accent, use_ansi);
42
43    if use_ansi {
44        print!("{}", WRAP_OFF);
45        io::stdout().flush()?;
46    }
47    for i in 0..art.len().max(info.len()) {
48        let art_line = art.get(i).map(String::as_str).unwrap_or(&blank_art);
49        let info_line = info.get(i).map(String::as_str).unwrap_or(EMPTY);
50
51        println!(" {} {}", art_line, info_line);
52    }
53    if use_ansi {
54        print!("{}", WRAP_ON);
55        io::stdout().flush()?;
56    }
57
58    println!();
59    Ok(())
60}
61
62/// Returns a vector of formatted strings representing the device info,
63/// including extra vendor-specific and general device limits.
64/// Lines for optional fields are only included if available.
65fn get_device_info(device: &Device, color: &str, use_ansi: bool) -> Vec<String> {
66    let mut lines = Vec::new();
67    let bold = if use_ansi { BOLD } else { EMPTY };
68    let reset = if use_ansi { RESET } else { EMPTY };
69    let value_color = if use_ansi { "\x1B[37m" } else { EMPTY };
70
71    let title = format!(
72        "{}{}{}{}: {}",
73        bold,
74        color,
75        device.device_name,
76        reset,
77        device.device_type.name()
78    );
79    let underline_len = device.device_name.len() + device.device_type.name().len() + 3;
80    let underline = "=".repeat(underline_len);
81
82    // Basic device info.
83    lines.push(title);
84    lines.push(format!("{}{}{}{}", bold, color, underline, reset));
85    lines.push(format!(
86        "{}{}Device{}: {}0x{:X}{} : {}0x{:X}{} ({})",
87        ALIGNMENT,
88        color,
89        reset,
90        value_color,
91        device.device_id,
92        reset,
93        value_color,
94        device.vendor_id,
95        reset,
96        device.vendor.name(),
97    ));
98    push_driver_info(&mut lines, device, color, value_color, reset);
99    lines.push(format!(
100        "{}{}API{}: {}{}{}",
101        ALIGNMENT, color, reset, value_color, device.api_version, reset
102    ));
103
104    let pressure_color = memory_pressure_color(device.characteristics.memory_pressure, use_ansi);
105    let heap_budget = device
106        .heapbudget
107        .map(format_bytes)
108        .unwrap_or_else(|| "???".to_string());
109    lines.push(format!(
110        "{}{}VRAM{}: {}{}{} / {}",
111        ALIGNMENT,
112        color,
113        reset,
114        pressure_color,
115        heap_budget,
116        reset,
117        format_bytes(device.heapsize)
118    ));
119
120    let pressure = device.characteristics.memory_pressure;
121    let pressure_text = format_memory_pressure(pressure);
122    lines.push(format!(
123        "{}{} % {}{}{}",
124        ALIGNMENT,
125        format_meter(30, pressure, use_ansi),
126        pressure_color,
127        pressure_text,
128        reset
129    ));
130
131    // Vendor-specific extra info.
132    if let Some(cu) = device.characteristics.compute_units {
133        let compute_units = match device.characteristics.active_compute_units {
134            Some(active) if active > 0 => format!("{} / {}", active, cu),
135            _ => cu.to_string(),
136        };
137        lines.push(format!(
138            "{}{}Compute Units{}: {}{}{}",
139            ALIGNMENT, color, reset, value_color, compute_units, reset
140        ));
141    }
142    if let Some(se) = device.characteristics.shader_engines {
143        lines.push(format!(
144            "{}{}Shader Engines{}: {}{}{}",
145            ALIGNMENT, color, reset, value_color, se, reset
146        ));
147    }
148    if let Some(sapec) = device.characteristics.shader_arrays_per_engine_count {
149        lines.push(format!(
150            "{}{}Shader Arrays per Engine{}: {}{}{}",
151            ALIGNMENT, color, reset, value_color, sapec, reset
152        ));
153    }
154    if let Some(cups) = device.characteristics.compute_units_per_shader_array {
155        lines.push(format!(
156            "{}{}Compute Units per Shader Array{}: {}{}{}",
157            ALIGNMENT, color, reset, value_color, cups, reset
158        ));
159    }
160    if let Some(simd) = device.characteristics.simd_per_compute_unit {
161        lines.push(format!(
162            "{}{}SIMD per Compute Unit{}: {}{}{}",
163            ALIGNMENT, color, reset, value_color, simd, reset
164        ));
165    }
166    if let Some(wfs) = device.characteristics.wavefronts_per_simd {
167        lines.push(format!(
168            "{}{}Wavefronts per SIMD{}: {}{}{}",
169            ALIGNMENT, color, reset, value_color, wfs, reset
170        ));
171    }
172    if let Some(wfsz) = device.characteristics.wavefront_size {
173        lines.push(format!(
174            "{}{}Wavefront Size{}: {}{}{}",
175            ALIGNMENT, color, reset, value_color, wfsz, reset
176        ));
177    }
178    if let Some(sm) = device.characteristics.streaming_multiprocessors {
179        lines.push(format!(
180            "{}{}Streaming Multiprocessors{}: {}{}{}",
181            ALIGNMENT, color, reset, value_color, sm, reset
182        ));
183    }
184    if let Some(wps) = device.characteristics.warps_per_sm {
185        lines.push(format!(
186            "{}{}Warps per SM{}: {}{}{}",
187            ALIGNMENT, color, reset, value_color, wps, reset
188        ));
189    }
190
191    // General device limits.
192    // lines.push(format!(
193    //     "{}{}Max Image Dimension 2D{}: {}",
194    //     ALIGNMENT,
195    //     color,
196    //     RESET,
197    //     format_bytes(device.characteristics.max_image_dimension_2d.into())
198    // ));
199    lines.push(format!(
200        "{}{}Max Compute Shared Memory Size{}: {}",
201        ALIGNMENT,
202        color,
203        reset,
204        format_bytes(device.characteristics.max_compute_shared_memory_size.into())
205    ));
206    lines.push(format!(
207        "{}{}Max Compute Work Group Invocations{}: {}",
208        ALIGNMENT, color, reset, device.characteristics.max_compute_work_group_invocations
209    ));
210
211    let checkbox = |b: bool| if b { "[x]" } else { "[ ]" };
212    let x = checkbox(device.characteristics.supports_ray_tracing);
213    let y = checkbox(device.characteristics.dedicated_transfer_queue);
214    let z = checkbox(device.characteristics.dedicated_async_compute_queue);
215
216    lines.push(format!(
217        "{}{}Raytracing{}: {} | {}Dedicated Transfer Queue{}: {} | {}Dedicated Async Compute Queue{}: {}",
218        ALIGNMENT,
219        color, reset, x,
220        color, reset, y,
221        color, reset, z,
222    ));
223
224    lines
225}
226
227fn push_driver_info(
228    lines: &mut Vec<String>,
229    device: &Device,
230    color: &str,
231    value_color: &str,
232    reset: &str,
233) {
234    let mut driver_info_lines = device.driver_info.lines().filter(|line| !line.is_empty());
235    match driver_info_lines.next() {
236        Some(first_line) => lines.push(format!(
237            "{}{}Driver{}: {}{}{} | {}{}{}",
238            ALIGNMENT,
239            color,
240            reset,
241            value_color,
242            device.driver_name,
243            reset,
244            value_color,
245            first_line,
246            reset
247        )),
248        None => lines.push(format!(
249            "{}{}Driver{}: {}{}{}",
250            ALIGNMENT, color, reset, value_color, device.driver_name, reset
251        )),
252    }
253
254    for line in driver_info_lines {
255        lines.push(format!("           {}{}{}", value_color, line, reset));
256    }
257}
258
259fn memory_pressure_color(pressure: Option<f32>, use_ansi: bool) -> &'static str {
260    if !use_ansi {
261        return EMPTY;
262    }
263
264    match pressure {
265        Some(pressure) if pressure < 0.5 => BRIGHT_GREEN,
266        Some(pressure) if pressure < 0.75 => BRIGHT_YELLOW,
267        Some(_) => BRIGHT_RED,
268        None => DIM,
269    }
270}
271
272fn format_memory_pressure(pressure: Option<f32>) -> String {
273    pressure
274        .filter(|pressure| pressure.is_finite())
275        .map(|pressure| format!("{:.2}", pressure * 100.0))
276        .unwrap_or_else(|| "???".to_string())
277}
278
279fn format_meter(width: usize, completion: Option<f32>, use_ansi: bool) -> String {
280    let inner_width = width.saturating_sub(2).max(1);
281    let completion = completion.filter(|completion| completion.is_finite() && *completion >= 0.0);
282    let mut result = String::with_capacity(width);
283
284    result.push('[');
285    for index in 0..inner_width {
286        match completion {
287            Some(completion) => {
288                let denominator = inner_width.saturating_sub(1).max(1) as f32;
289                let phase = index as f32 / denominator;
290                if phase <= completion.clamp(0.0, 1.0) {
291                    result.push_str(memory_pressure_color(Some(phase), use_ansi));
292                    result.push('|');
293                } else {
294                    result.push(' ');
295                }
296            }
297            None => {
298                if use_ansi {
299                    result.push_str(DIM);
300                }
301                result.push('-');
302            }
303        }
304    }
305    if use_ansi {
306        result.push_str(RESET);
307    }
308    result.push(']');
309    result
310}
311
312/// Converts a byte count into a human-readable string with binary units.
313fn format_bytes(bytes: u64) -> String {
314    const UNITS: [&str; 9] = [
315        "Bytes", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB",
316    ];
317
318    let mut size = bytes as f64;
319    let mut unit = 0;
320    while size >= 1024.0 && unit + 1 < UNITS.len() {
321        size /= 1024.0;
322        unit += 1;
323    }
324
325    format!("{:.3} {}", size, UNITS[unit])
326}
327
328/// Iterates through API versions and prints info for every physical device
329pub fn iterate_devices() -> Result<(), Box<dyn Error>> {
330    let entry = {
331        #[cfg(not(feature = "loaded"))]
332        {
333            Entry::linked()
334        }
335        #[cfg(feature = "loaded")]
336        {
337            unsafe { Entry::load().map_err(|err| io::Error::other(format!("{err:?}")))? }
338        }
339    };
340
341    let mut last_create_error = None;
342    for api_version in [
343        vk::API_VERSION_1_3,
344        vk::API_VERSION_1_2,
345        vk::API_VERSION_1_1,
346        vk::API_VERSION_1_0,
347    ] {
348        let app_info = vk::ApplicationInfo::default()
349            .application_name(c"vkfetch-rs")
350            .application_version(package_version())
351            .engine_name(c"vkfetch-rs")
352            .engine_version(package_version())
353            .api_version(api_version);
354
355        let mut extension_names = Vec::new();
356        let mut flags = vk::InstanceCreateFlags::empty();
357        if supports_instance_extension(&entry, vk::KHR_PORTABILITY_ENUMERATION_NAME) {
358            extension_names.push(vk::KHR_PORTABILITY_ENUMERATION_NAME.as_ptr());
359            flags |= vk::InstanceCreateFlags::ENUMERATE_PORTABILITY_KHR;
360        }
361
362        let create_info = vk::InstanceCreateInfo::default()
363            .application_info(&app_info)
364            .enabled_extension_names(&extension_names)
365            .flags(flags);
366
367        match unsafe { entry.create_instance(&create_info, None) } {
368            Ok(instance) => {
369                match unsafe { instance.enumerate_physical_devices() } {
370                    Ok(devices) => {
371                        for device in devices {
372                            if let Err(error) = fetch_device(&instance, device) {
373                                unsafe {
374                                    instance.destroy_instance(None);
375                                }
376                                return Err(error);
377                            }
378                        }
379                    }
380                    Err(e) => {
381                        eprintln!("Failed to enumerate physical devices: {:?}", e);
382                    }
383                };
384                unsafe {
385                    instance.destroy_instance(None);
386                }
387            }
388            Err(e) => {
389                eprintln!("Failed to create instance: {:?}", e);
390                last_create_error = Some(e);
391                continue;
392            }
393        };
394
395        return Ok(());
396    }
397
398    match last_create_error {
399        Some(error) => {
400            Err(io::Error::other(format!("failed to create Vulkan instance: {error:?}")).into())
401        }
402        None => Ok(()),
403    }
404}
405
406fn supports_instance_extension(entry: &Entry, extension_name: &CStr) -> bool {
407    let Ok(extensions) = (unsafe { entry.enumerate_instance_extension_properties(None) }) else {
408        return false;
409    };
410
411    extensions.iter().any(|extension| {
412        extension
413            .extension_name_as_c_str()
414            .is_ok_and(|name| name == extension_name)
415    })
416}
417
418fn package_version() -> u32 {
419    let major = env!("CARGO_PKG_VERSION_MAJOR").parse().unwrap_or(0);
420    let minor = env!("CARGO_PKG_VERSION_MINOR").parse().unwrap_or(0);
421    let patch = env!("CARGO_PKG_VERSION_PATCH").parse().unwrap_or(0);
422    vk::make_api_version(0, major, minor, patch)
423}
424
425#[cfg(windows)]
426mod vt {
427    use std::io::{Error, Result};
428    use winapi::um::consoleapi::{GetConsoleMode, SetConsoleMode};
429    use winapi::um::handleapi::INVALID_HANDLE_VALUE;
430    use winapi::um::processenv::GetStdHandle;
431    use winapi::um::winbase::STD_OUTPUT_HANDLE;
432    use winapi::um::wincon::ENABLE_VIRTUAL_TERMINAL_PROCESSING;
433
434    /// Enables Virtual Terminal Processing on Windows.
435    pub fn enable_virtual_terminal_processing() -> Result<()> {
436        unsafe {
437            let std_out = GetStdHandle(STD_OUTPUT_HANDLE);
438            if std_out == INVALID_HANDLE_VALUE {
439                return Err(Error::last_os_error());
440            }
441            let mut mode = 0;
442            if GetConsoleMode(std_out, &mut mode) == 0 {
443                return Err(Error::last_os_error());
444            }
445            mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
446            if SetConsoleMode(std_out, mode) == 0 {
447                return Err(Error::last_os_error());
448            }
449        }
450        Ok(())
451    }
452
453    /// Checks if Virtual Terminal Processing is enabled.
454    pub fn is_vt_enabled() -> bool {
455        unsafe {
456            let std_out = GetStdHandle(STD_OUTPUT_HANDLE);
457            if std_out == INVALID_HANDLE_VALUE {
458                return false;
459            }
460            let mut mode = 0;
461            if GetConsoleMode(std_out, &mut mode) == 0 {
462                return false;
463            }
464            (mode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) != 0
465        }
466    }
467}
468
469#[cfg(not(windows))]
470mod vt {
471    use std::io::Result;
472
473    /// On non-Windows platforms, VT processing is typically enabled by default.
474    pub fn enable_virtual_terminal_processing() -> Result<()> {
475        Ok(())
476    }
477
478    /// Assume ANSI escape codes are supported.
479    pub fn is_vt_enabled() -> bool {
480        true
481    }
482}
483
484/// Returns `true` if stdout is a TTY and (on Windows) VT processing is enabled.
485fn is_ansi_supported() -> bool {
486    std::io::IsTerminal::is_terminal(&std::io::stdout()) && vt::is_vt_enabled()
487}
488
489#[cfg(test)]
490mod tests {
491    use super::*;
492    use crate::device::{Device, GPUCharacteristics};
493    use crate::vendor::Vendor;
494
495    /// For testing purposes we use the Unknown vendor variant.
496    impl Vendor {
497        pub fn dummy() -> Self {
498            Vendor::Unknown
499        }
500    }
501
502    /// Creates a dummy PhysicalDevice instance for tests.
503    fn dummy_physical_device() -> Device {
504        Device {
505            vendor: Vendor::dummy(),
506            device_name: "TestDevice".to_string(),
507            device_type: crate::device::DeviceType::DiscreteGPU,
508            device_id: 0xDEADBEEF,
509            vendor_id: 0xBEEF,
510            driver_name: "TestDriver".to_string(),
511            driver_info: "TestDriverInfo\nSecond line".to_string(),
512            api_version: "1.2.3.4".to_string(),
513            heapbudget: Some(8 * 1024 * 1024 * 1024), // 8 GiB
514            heapsize: 10 * 1024 * 1024 * 1024,        // 10 GB
515            characteristics: GPUCharacteristics {
516                memory_pressure: Some(0.2), // 20%
517                compute_units: Some(10),
518                active_compute_units: Some(8),
519                shader_engines: Some(2),
520                shader_arrays_per_engine_count: Some(2),
521                compute_units_per_shader_array: Some(5),
522                simd_per_compute_unit: Some(64),
523                wavefronts_per_simd: Some(4),
524                wavefront_size: Some(32),
525                streaming_multiprocessors: Some(46),
526                warps_per_sm: Some(32),
527                max_image_dimension_2d: 16384,
528                max_compute_shared_memory_size: 65536,
529                max_compute_work_group_invocations: 1024,
530                dedicated_transfer_queue: true,
531                dedicated_async_compute_queue: true,
532                supports_ray_tracing: true,
533            },
534        }
535    }
536
537    #[test]
538    fn test_format_bytes() {
539        assert_eq!(format_bytes(500), "500.000 Bytes");
540        assert_eq!(format_bytes(1024), "1.000 KiB");
541        assert_eq!(format_bytes(1024 * 1024), "1.000 MiB");
542        assert_eq!(format_bytes(1024 * 1024 * 1024), "1.000 GiB");
543        assert_eq!(format_bytes(1024 * 1024 * 1024 * 1024), "1.000 TiB");
544    }
545
546    #[test]
547    fn test_get_device_info() {
548        let device = dummy_physical_device();
549        let color = "\x1B[32m";
550        let info = get_device_info(&device, color, true);
551        assert!(info.len() >= 9);
552        assert!(info[0].contains("TestDevice"));
553        assert!(info[0].contains(device.device_type.name()));
554        assert!(info[2].contains("0xDEADBEEF"));
555        assert!(info[2].contains("0xBEEF"));
556        assert!(info.iter().any(|line| line.contains("Second line")));
557        assert!(info.iter().any(|line| line.contains("8 / 10")));
558        assert!(info.iter().any(|line| line.contains("32")));
559    }
560
561    #[test]
562    fn test_get_device_info_without_ansi() {
563        let device = dummy_physical_device();
564        let info = get_device_info(&device, EMPTY, false);
565        assert!(!info.iter().any(|line| line.contains("\x1B")));
566        assert!(info.iter().any(|line| line.contains("8.000 GiB")));
567    }
568
569    #[test]
570    fn test_unknown_memory_pressure_meter() {
571        let meter = format_meter(6, None, false);
572        assert_eq!(meter, "[----]");
573    }
574}
vkfetch_rs/lib.rs

vkfetch_rs/
lib.rs
