Skip to content

Cooperative kernel launch #508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
abionics wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Cooperative kernel launch #508

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6674bf9
Add cooperative kernel launch
abionics Oct 31, 2025
3f92d19
Support cooperative kernel launches on CUDA 9+
abionics Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 30 additions & 4 deletions doc/driver.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,12 @@ Constants

CUDA 8.0 and above.

.. attribute :: COOPERATIVE_LAUNCH

Whether the device supports cooperative kernel launches.

CUDA 9.0 and above.

.. attribute :: MAX_SHARED_MEMORY_PER_BLOCK_OPTIN

CUDA 9.0 and above.
Expand Down Expand Up @@ -1811,7 +1817,7 @@ Code on the Device: Modules and Functions
Handle to a *__global__* function in a :class:`Module`. Create using
:meth:`Module.get_function`.

.. method:: __call__(arg1, ..., argn, block=block_size, [grid=(1,1), [stream=None, [shared=0, [texrefs=[], [time_kernel=False]]]]])
.. method:: __call__(arg1, ..., argn, block=block_size, [grid=(1,1), [stream=None, [shared=0, [texrefs=[], [time_kernel=False, [cooperative=False]]]]]])

Launch *self*, with a thread block size of *block*. *block* must be a 3-tuple
of integers.
Expand All @@ -1828,6 +1834,13 @@ Code on the Device: Modules and Functions
*extern __shared__* arrays.
*texrefs* is a :class:`list` of :class:`TextureReference` instances
that the function will have access to.
If *cooperative* is *True*, launch the kernel cooperatively. This requires
CUDA 9.0 or newer and a device whose :attr:`device_attribute.COOPERATIVE_LAUNCH`
attribute is nonzero. The entire grid must fit concurrently on the device;
:meth:`get_max_active_blocks_per_multiprocessor` can be used to determine
this limit. On CUDA versions before 11.0, kernels that synchronize the
entire grid must be built using ``pycuda.compiler.DynamicSourceModule``.
Cooperatively launched kernels may not use CUDA dynamic parallelism.

The function returns either *None* or the number of seconds spent
executing the kernel, depending on whether *time_kernel* is *True*.
Expand Down Expand Up @@ -1893,18 +1906,20 @@ Code on the Device: Modules and Functions

Return `self`.

.. method:: prepared_call(grid, block, *args, shared_size=0)
.. method:: prepared_call(grid, block, *args, shared_size=0, cooperative=False)

Invoke `self` using :meth:`launch_grid`, with `args` a grid size of `grid`,
and a block size of *block*.
Assumes that :meth:`prepare` was called on *self*.
The texture references given to :meth:`prepare` are set up as parameters, as
well.
If *cooperative* is *True*, use the cooperative launch described in
:meth:`__call__`.

.. versionchanged:: 2012.1
*shared_size* was added.

.. method:: prepared_timed_call(grid, block, *args, shared_size=0)
.. method:: prepared_timed_call(grid, block, *args, shared_size=0, cooperative=False)

Invoke `self` using :meth:`launch_grid`, with `args`, a grid size of `grid`,
and a block size of *block*.
Expand All @@ -1919,7 +1934,7 @@ Code on the Device: Modules and Functions
.. versionchanged:: 2012.1
*shared_size* was added.

.. method:: prepared_async_call(grid, block, stream, *args, shared_size=0)
.. method:: prepared_async_call(grid, block, stream, *args, shared_size=0, cooperative=False)

Invoke `self` using :meth:`launch_grid_async`, with `args`, a grid size
of `grid`, and a block size of *block*, serialized into the
Expand All @@ -1931,6 +1946,17 @@ Code on the Device: Modules and Functions
.. versionchanged:: 2012.1
*shared_size* was added.

.. method:: get_max_active_blocks_per_multiprocessor(block_size, dynamic_smem_size=0)

Return the maximum number of thread blocks for *self* that can be active
on one multiprocessor. *dynamic_smem_size* is the dynamic shared memory
used by each block, in bytes.

A cooperative launch may contain at most this value multiplied by the
device's :attr:`device_attribute.MULTIPROCESSOR_COUNT`.

CUDA 9.0 and above.

.. method:: get_attribute(attr)

Return one of the attributes given by the
Expand Down
85 changes: 53 additions & 32 deletions pycuda/driver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,37 +197,37 @@ def device_get_attributes(dev):
def device___getattr__(dev, name):
return dev.get_attribute(getattr(device_attribute, name.upper()))

def _build_arg_buf(args):
def _build_arg_buf(args, cooperative=False):
handlers = []

arg_data = []
format = ""
formats = []
for i, arg in enumerate(args):
if isinstance(arg, np.number):
arg_data.append(arg)
format += arg.dtype.char
formats.append(arg.dtype.char)
elif isinstance(arg, (DeviceAllocation, PooledDeviceAllocation)):
arg_data.append(int(arg))
format += "P"
formats.append("P")
elif isinstance(arg, ArgumentHandler):
handlers.append(arg)
arg_data.append(int(arg.get_device_alloc()))
format += "P"
formats.append("P")
elif isinstance(arg, np.ndarray):
if isinstance(arg.base, ManagedAllocationOrStub):
arg_data.append(int(arg.base))
format += "P"
formats.append("P")
else:
arg_data.append(arg)
format += "%ds" % arg.nbytes
formats.append("%ds" % arg.nbytes)
elif isinstance(arg, np.void):
arg_data.append(_my_bytes(_memoryview(arg)))
format += "%ds" % arg.itemsize
formats.append("%ds" % arg.itemsize)
else:
cai = getattr(arg, "__cuda_array_interface__", None)
if cai:
arg_data.append(cai["data"][0])
format += "P"
formats.append("P")
continue

try:
Expand All @@ -237,11 +237,16 @@ def _build_arg_buf(args):
else:
# for gpuarrays
arg_data.append(int(gpudata))
format += "P"
formats.append("P")

from pycuda._pvt_struct import pack

return handlers, pack(format, *arg_data)
if cooperative:
return handlers, [
pack(formats[i], arg) for i, arg in enumerate(arg_data)
]

return handlers, pack("".join(formats), *arg_data)

# {{{ pre-CUDA 4 call interface (stateful)

Expand Down Expand Up @@ -481,6 +486,7 @@ def function_call(func, *args, **kwargs):
shared = kwargs.pop("shared", 0)
texrefs = kwargs.pop("texrefs", [])
time_kernel = kwargs.pop("time_kernel", False)
cooperative = kwargs.pop("cooperative", False)

if kwargs:
raise ValueError(
Expand All @@ -491,7 +497,7 @@ def function_call(func, *args, **kwargs):
raise ValueError("must specify block size")

func._set_block_shape(*block)
handlers, arg_buf = _build_arg_buf(args)
handlers, arg_buf = _build_arg_buf(args, cooperative)

for handler in handlers:
handler.pre_call(stream)
Expand All @@ -511,7 +517,7 @@ def function_call(func, *args, **kwargs):

start_time = time()

func._launch_kernel(grid, block, arg_buf, shared, None)
func._launch_kernel(grid, block, arg_buf, shared, None, cooperative)

if post_handlers or time_kernel:
Context.synchronize()
Expand All @@ -528,7 +534,7 @@ def function_call(func, *args, **kwargs):
assert (
not time_kernel
), "Can't time the kernel on an asynchronous invocation"
func._launch_kernel(grid, block, arg_buf, shared, stream)
func._launch_kernel(grid, block, arg_buf, shared, stream, cooperative)

if post_handlers:
for handler in post_handlers:
Expand All @@ -539,23 +545,41 @@ def function_prepare(func, arg_types, texrefs=None):
texrefs = []
func.texrefs = texrefs

func.arg_format = ""
func.arg_formats = []

for _i, arg_type in enumerate(arg_types):
if isinstance(arg_type, type) and np.number in arg_type.__mro__:
func.arg_format += np.dtype(arg_type).char
arg_format = np.dtype(arg_type).char
elif isinstance(arg_type, np.dtype):
if arg_type.char == "V":
func.arg_format += "%ds" % arg_type.itemsize
arg_format = "%ds" % arg_type.itemsize
else:
func.arg_format += arg_type.char
arg_format = arg_type.char
elif isinstance(arg_type, str):
func.arg_format += arg_type
arg_format = arg_type
else:
func.arg_format += np.dtype(np.uintp).char
arg_format = np.dtype(np.uintp).char

func.arg_formats.append(arg_format)

func.arg_format = "".join(func.arg_formats)

return func

def _build_prepared_arg_buf(func, args, cooperative):
from pycuda._pvt_struct import pack

if not cooperative:
return pack(func.arg_format, *args)

if len(args) != len(func.arg_formats):
raise TypeError(
"expected %d kernel arguments, got %d"
% (len(func.arg_formats), len(args))
)

return [pack(func.arg_formats[i], arg) for i, arg in enumerate(args)]

def function_prepared_call(func, grid, block, *args, **kwargs):
if isinstance(block, tuple):
func._set_block_shape(*block)
Expand All @@ -571,31 +595,29 @@ def function_prepared_call(func, grid, block, *args, **kwargs):
args = (block, *args)

shared_size = kwargs.pop("shared_size", 0)
cooperative = kwargs.pop("cooperative", False)

if kwargs:
raise TypeError(
"unknown keyword arguments: " + ", ".join(kwargs.keys())
)

from pycuda._pvt_struct import pack

arg_buf = pack(func.arg_format, *args)
arg_buf = _build_prepared_arg_buf(func, args, cooperative)

for texref in func.texrefs:
func.param_set_texref(texref)

func._launch_kernel(grid, block, arg_buf, shared_size, None)
func._launch_kernel(grid, block, arg_buf, shared_size, None, cooperative)

def function_prepared_timed_call(func, grid, block, *args, **kwargs):
shared_size = kwargs.pop("shared_size", 0)
cooperative = kwargs.pop("cooperative", False)
if kwargs:
raise TypeError(
"unknown keyword arguments: " + ", ".join(kwargs.keys())
)

from pycuda._pvt_struct import pack

arg_buf = pack(func.arg_format, *args)
arg_buf = _build_prepared_arg_buf(func, args, cooperative)

for texref in func.texrefs:
func.param_set_texref(texref)
Expand All @@ -604,7 +626,7 @@ def function_prepared_timed_call(func, grid, block, *args, **kwargs):
end = Event()

start.record()
func._launch_kernel(grid, block, arg_buf, shared_size, None)
func._launch_kernel(grid, block, arg_buf, shared_size, None, cooperative)
end.record()

def get_call_time():
Expand All @@ -629,20 +651,19 @@ def function_prepared_async_call(func, grid, block, stream, *args, **kwargs):
stream = block

shared_size = kwargs.pop("shared_size", 0)
cooperative = kwargs.pop("cooperative", False)

if kwargs:
raise TypeError(
"unknown keyword arguments: " + ", ".join(kwargs.keys())
)

from pycuda._pvt_struct import pack

arg_buf = pack(func.arg_format, *args)
arg_buf = _build_prepared_arg_buf(func, args, cooperative)

for texref in func.texrefs:
func.param_set_texref(texref)

func._launch_kernel(grid, block, arg_buf, shared_size, stream)
func._launch_kernel(grid, block, arg_buf, shared_size, stream, cooperative)

# }}}

Expand Down
Loading
Loading