Use NVTX.jl to give streams a name

[simonbyrne]

We could add CUDA as a WeakDep?

[simonbyrne]

See this commented out block: https://github.com/JuliaGPU/CUDA.jl/blob/2ae53a46658f10b0be32ae61fb4e145782ed9208/lib/cudadrv/state.jl#L368-L373

Unfortunately these are not included in our JLL binaries.

[simonbyrne]

@maleadt do you know if nvtxNameCuStreamA is available in any of the CUDA packages? If not, we may need to build the NVTX library differently.

[simonbyrne]

Ok, these are exposed in the new JLLs. @maleadt do you have a preference if NVTX.jl should have CUDA.jl as a weak dep, or vice versa?

For an interface, I was thinking of adding a function along the following lines:

NVTX.name_resource(stream::CuStream, name::AbstractString)
NVTX.name_resource(stream::CuContext, name::AbstractString)
NVTX.name_resource(stream::CuDevice, name::AbstractString)

[maleadt]

Adding NVTX.jl as a weak dep to CUDA.jl seems like the more logical option to me, as it's CUDA.jl that needs to name streams when they get created. Or how would you implement the inverse?

[simonbyrne]

I think you would need users to manually name their own streams?

[simonbyrne]

Adding NVTX.jl as a weak dep to CUDA.jl seems like the more logical option to me, as it's CUDA.jl that needs to name streams when they get created.

Alternatively: you could just add NVTX_jll as a dependency to CUDA, and directly ccall it there, similar to what is currently commented out: https://github.com/JuliaGPU/CUDA.jl/blob/2ae53a46658f10b0be32ae61fb4e145782ed9208/lib/cudadrv/state.jl#L368-L373

We could also then define the above functions here (with CUDA as a weak dep) for users who want to give more descriptive names to their streams?

[maleadt]

I think you would need users to manually name their own streams?

Why? We can have CUDA at least provide a somewhat better name on stream construction, and the user can always improve afterwards.

Alternatively: you could just add NVTX_jll as a dependency to CUDA, and directly ccall it there, similar to what is currently commented out: https://github.com/JuliaGPU/CUDA.jl/blob/2ae53a46658f10b0be32ae61fb4e145782ed9208/lib/cudadrv/state.jl#L368-L373

Why can't we do that through a weakdep? If the user isn't using NVTX.jl, I assume he wouldn't be interested in user-friendly names for CUDA.jl's streams, no?

[simonbyrne]

Can you call weakdeps from inside functions in the package? Or how would it work?

[maleadt]

Add name_stream(x) = nothing to CUDA.jl and have the weakdep provide a more specific name_stream(x::CuStream) = NVTX....? Or do something dynamic a la https://github.com/JuliaGPU/GPUCompiler.jl/blob/661dcef51a93f4ce76cb9446d7d8670cd97ea8ca/src/reflection.jl#L65-L69. @KristofferC may know what the best option is here (summary: we want CUDA.jl to call functionality from NVTX.jl when the package is available, using weakdeps so that we can express compatibility constraints).

[KristofferC]

(summary: we want CUDA.jl to call functionality from NVTX.jl when the package is available, using weakdeps so that we can express compatibility constraints).

If by "available" you mean that the package is loaded, you have two options:

1. Overload some method in CUDA.jl in the extension. For example, have a default definition stream_name(x) = "default name and in the extension have stream_name(x::CudaStream) = NVTX.cool_stream_name() and always call it with a CudaStream.
2. Do a dynamic check if the extension is loaded using Base.get_extension.

Extensions work best if the user provides something (like a type) in the weak dependency to dispatch on but that doesn't seem to be the case here.

[simonbyrne]

Since CUDA.jl now depends on NVTX.jl, how about the following:

• add a function stub (name_resource?) to NVTX.jjl
• add the methods in CUDA.jl

[maleadt]

Wouldn't it be simpler to add a high-level API for nvtxNameCuStream to NVTX.jl and have CUDA.jl call that on stream construction?

[simonbyrne]

CuStream is defined in CUDA.jl. We could do it via package extensions, but since CUDA.jl loads NVTX.jl, it seems like asker to do it here.