nvtx 2.0.0

NVIDIA Tools Extension SDK
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

use crate::Str;

/// Enum for all CUDA Runtime types.
pub enum CudaRuntimeResource {
    /// Device (integer ID)
    Device(i32),
    /// Event
    Event(nvtx_sys::CudaEvent),
    /// Stream
    Stream(nvtx_sys::CudaStream),
}

impl From<i32> for CudaRuntimeResource {
    fn from(value: i32) -> Self {
        CudaRuntimeResource::Device(value)
    }
}

impl From<nvtx_sys::CudaEvent> for CudaRuntimeResource {
    fn from(value: nvtx_sys::CudaEvent) -> Self {
        CudaRuntimeResource::Event(value)
    }
}

impl From<nvtx_sys::CudaStream> for CudaRuntimeResource {
    fn from(value: nvtx_sys::CudaStream) -> Self {
        CudaRuntimeResource::Stream(value)
    }
}

/// Name a CUDA Runtime Resource (one of: Device, Event, or Stream).
///
/// ```
/// nvtx::name_cudart_resource(nvtx::CudaRuntimeResource::Device(0), c"GPU 0");
/// /// or implicitly:
/// nvtx::name_cudart_resource(0, c"GPU 0");
/// ```
pub fn name_cudart_resource(resource: impl Into<CudaRuntimeResource>, name: impl Into<Str>) {
    match resource.into() {
        CudaRuntimeResource::Device(device) => match name.into() {
            Str::Ascii(s) => nvtx_sys::name_cuda_device_ascii(device, &s),
            Str::Unicode(s) => nvtx_sys::name_cuda_device_unicode(device, &s),
        },
        CudaRuntimeResource::Event(event) => match name.into() {
            Str::Ascii(s) => {
                // SAFETY: NVTX requires a valid CUDA runtime event handle; caller provides `CudaRuntimeResource::Event`.
                unsafe { nvtx_sys::name_cuda_event_ascii(event, &s) }
            }
            Str::Unicode(s) => {
                // SAFETY: NVTX requires a valid CUDA runtime event handle; caller provides `CudaRuntimeResource::Event`.
                unsafe { nvtx_sys::name_cuda_event_unicode(event, &s) }
            }
        },
        CudaRuntimeResource::Stream(stream) => match name.into() {
            Str::Ascii(s) => {
                // SAFETY: NVTX requires a valid CUDA runtime stream handle; caller provides `CudaRuntimeResource::Stream`.
                unsafe { nvtx_sys::name_cuda_stream_ascii(stream, &s) }
            }
            Str::Unicode(s) => {
                // SAFETY: NVTX requires a valid CUDA runtime stream handle; caller provides `CudaRuntimeResource::Stream`.
                unsafe { nvtx_sys::name_cuda_stream_unicode(stream, &s) }
            }
        },
    }
}