[abionics]

Summary

Add cooperative kernel launches to PyCUDA through cooperative=True, using cuLaunchCooperativeKernel on CUDA 9+.

The option is supported by direct, prepared, timed, and asynchronous calls while preserving existing non-cooperative behavior.

func(                                                                                                                                                       
    buf_gpu,                                                                                                                                                
    block=(256, 1, 1),                                                                                                                                      
    grid=(2, 1, 1),                                                                                                                                         
    cooperative=True,                                                                                                                                       
)                                                                                                                                                           

Prepared functions use the same option:

func.prepare("P")                                                                                                                                           
func.prepared_call(                                                                                                                                         
    (2, 1, 1),                                                                                                                                              
    (256, 1, 1),                                                                                                                                            
    buf_gpu,                                                                                                                                                
    cooperative=True,                                                                                                                                       
)                                                                                                                                                           

A cooperative-groups kernel can be compiled normally on CUDA 11+:

from pycuda.compiler import SourceModule                                                                                                                    

mod = SourceModule(r"""
#include <cooperative_groups.h>

extern "C" __global__ void kernel(float *buf)
{
    cooperative_groups::grid_group grid =
        cooperative_groups::this_grid();

    grid.sync();
}
""", no_extern_c=True)

On CUDA 9_10, use DynamicSourceModule for kernels performing grid-wide synchronization.

Additional changes

• Correct per-argument packing for cooperative launches.
• Expose device_attribute.COOPERATIVE_LAUNCH.
• Add get_max_active_blocks_per_multiprocessor() for grid sizing.
• Report oversized cooperative grids as LaunchError.
• Add focused tests and documentation.

Tested on CUDA 13. All 22 driver tests passed.

Closes #253.

[inducer]

Thanks for making this happen! I'll try to take a look sometime this upcoming week.

[inducer]

Thanks for making this! Almost looks good, except for one compatibility oddity.

[inducer]

There's a bit of a compatibility issue here. Hypothetically, struct.pack formats that are more complex than a single character are supported. (Spaces are allowed, too.) This code supports that by accident, in a way, but we have to assume it's being used.

Cooperative wouldn't support this usage as-is. Here's a struct format parser I've just slopped together:

https://github.com/proxy/gist.github.com/inducer/07ac13dd1ac7827eb21dcd372986880f

That could help with this. (Needs to be extended for the complex values supported by PyCUDA.) Since this would happen during prep, I don't imagine the cost would be noticeable/prohibitive.