Running Hopper TMA kernels with pycuda

[bastianhagedorn]

Hi, I would like to use pycuda to compile and launch kernels targeting NVIDIAs Hopper architecture. These kernels might use the TMA hardware unit (https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/) which requires creating CUTensorMap objects (https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__TENSOR__MEMORY.html) on the host side and pass them as additional kernel parameters to the kernel.
Would it be possible to express this somehow in pycuda?

[inducer]

Not yet, but contributions that wrap these should be straightforward (just by following the existing patterns) and are definitely welcome.

[bastianhagedorn]

@inducer Can you point me to some of these existing patterns?
If it's that straightforward I'd also very much appreciate if someone could add this support because as a pure user of pycuda I'm completly unfamiliar with its code base.
IIRC in the past, you similarly quickly added support for setting dynamic shared memory as a kernel function attribute which I'm still using as of today :)

Anyways, any progress in supporting that would be greatly appreciated!

[bastianhagedorn]

@inducer pycuda uses nvrtc under the hood, correct?

[inducer]

There was a proposed change at a point, but that sort of petered out. Currently, we run nvcc in a shell under the hood.