diff --git a/doc/driver.rst b/doc/driver.rst index d1fb596a..f8875ea9 100644 --- a/doc/driver.rst +++ b/doc/driver.rst @@ -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. @@ -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. @@ -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*. @@ -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*. @@ -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 @@ -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 diff --git a/pycuda/driver.py b/pycuda/driver.py index 45822b64..4f2ff937 100644 --- a/pycuda/driver.py +++ b/pycuda/driver.py @@ -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: @@ -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) @@ -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( @@ -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) @@ -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() @@ -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: @@ -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) @@ -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) @@ -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(): @@ -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) # }}} diff --git a/src/cpp/cuda.hpp b/src/cpp/cuda.hpp index 685216de..484c5de1 100644 --- a/src/cpp/cuda.hpp +++ b/src/cpp/cuda.hpp @@ -1445,10 +1445,23 @@ namespace pycuda } #endif +#if CUDAPP_CUDA_VERSION >= 9000 + int get_max_active_blocks_per_multiprocessor( + int block_size, size_t dynamic_smem_size=0) const + { + int result; + CUDAPP_CALL_GUARDED_WITH_TRACE_INFO( + cuOccupancyMaxActiveBlocksPerMultiprocessor, + (&result, m_function, block_size, dynamic_smem_size), m_symbol); + return result; + } +#endif + #if CUDAPP_CUDA_VERSION >= 4000 void launch_kernel(py::tuple grid_dim_py, py::tuple block_dim_py, py::object parameter_buffer, - unsigned shared_mem_bytes, py::object stream_py) + unsigned shared_mem_bytes, py::object stream_py, + bool cooperative=false) { const unsigned axis_count = 3; unsigned grid_dim[axis_count]; @@ -1478,22 +1491,62 @@ namespace pycuda PYCUDA_PARSE_STREAM_PY; - py_buffer_wrapper par_buf_wrapper; - par_buf_wrapper.get(parameter_buffer.ptr(), PyBUF_ANY_CONTIGUOUS); - size_t par_len = par_buf_wrapper.m_buf.len; - - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, const_cast(par_buf_wrapper.m_buf.buf), - CU_LAUNCH_PARAM_BUFFER_SIZE, &par_len, - CU_LAUNCH_PARAM_END - }; - - CUDAPP_CALL_GUARDED( - cuLaunchKernel, (m_function, - grid_dim[0], grid_dim[1], grid_dim[2], - block_dim[0], block_dim[1], block_dim[2], - shared_mem_bytes, s_handle, 0, config - )); + if (!cooperative) + { + py_buffer_wrapper par_buf_wrapper; + par_buf_wrapper.get(parameter_buffer.ptr(), PyBUF_ANY_CONTIGUOUS); + size_t par_len = par_buf_wrapper.m_buf.len; + + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + const_cast(par_buf_wrapper.m_buf.buf), + CU_LAUNCH_PARAM_BUFFER_SIZE, &par_len, + CU_LAUNCH_PARAM_END + }; + + CUDAPP_CALL_GUARDED( + cuLaunchKernel, (m_function, + grid_dim[0], grid_dim[1], grid_dim[2], + block_dim[0], block_dim[1], block_dim[2], + shared_mem_bytes, s_handle, 0, config + )); + } + else + { +#if CUDAPP_CUDA_VERSION >= 9000 + pycuda_size_t param_count = py::len(parameter_buffer); + std::vector > param_wrappers; + std::vector kernel_params; + param_wrappers.reserve(param_count); + kernel_params.reserve(param_count); + + for (pycuda_size_t i = 0; i < param_count; ++i) + { + py::object param = parameter_buffer[i]; + param_wrappers.push_back( + std::unique_ptr(new py_buffer_wrapper())); + param_wrappers.back()->get(param.ptr(), PyBUF_ANY_CONTIGUOUS); + kernel_params.push_back(param_wrappers.back()->m_buf.buf); + } + + CUDAPP_CALL_GUARDED( + cuLaunchCooperativeKernel, (m_function, + grid_dim[0], grid_dim[1], grid_dim[2], + block_dim[0], block_dim[1], block_dim[2], + shared_mem_bytes, s_handle, + kernel_params.empty() ? 0 : &kernel_params.front() + )); +#else + throw pycuda::error("function::launch_kernel", +#if CUDAPP_CUDA_VERSION >= 5000 + CUDA_ERROR_NOT_SUPPORTED, +#else + CUDA_ERROR_INVALID_VALUE, +#endif + "cooperative kernel launch requires CUDA >= 9.0 " + "(use a newer CUDA toolkit, or avoid cooperative=True)."); +#endif + } } #endif diff --git a/src/wrapper/wrap_cudadrv.cpp b/src/wrapper/wrap_cudadrv.cpp index e881d684..b119ee5a 100644 --- a/src/wrapper/wrap_cudadrv.cpp +++ b/src/wrapper/wrap_cudadrv.cpp @@ -49,6 +49,9 @@ namespace || err.code() == CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES || err.code() == CUDA_ERROR_LAUNCH_TIMEOUT || err.code() == CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING +#if CUDAPP_CUDA_VERSION >= 9000 + || err.code() == CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE +#endif ) PyErr_SetString(CudaLaunchError.get(), err.what()); else if (err.code() == CUDA_ERROR_OUT_OF_MEMORY) @@ -910,6 +913,7 @@ BOOST_PYTHON_MODULE(_driver) .value("CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM", CU_DEVICE_ATTRIBUTE_CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM) #endif #if CUDAPP_CUDA_VERSION >= 9000 + .value("COOPERATIVE_LAUNCH", CU_DEVICE_ATTRIBUTE_COOPERATIVE_LAUNCH) .value("MAX_SHARED_MEMORY_PER_BLOCK_OPTIN", CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) #endif #if CUDAPP_CUDA_VERSION >= 9020 @@ -1323,11 +1327,19 @@ BOOST_PYTHON_MODULE(_driver) #if CUDAPP_CUDA_VERSION >= 10000 .DEF_SIMPLE_METHOD(set_attribute) #endif +#if CUDAPP_CUDA_VERSION >= 9000 + .def("get_max_active_blocks_per_multiprocessor", + &cl::get_max_active_blocks_per_multiprocessor, + (py::arg("block_size"), py::arg("dynamic_smem_size")=0)) +#endif #if CUDAPP_CUDA_VERSION >= 3000 && defined(CUDAPP_POST_30_BETA) .DEF_SIMPLE_METHOD(set_cache_config) #endif #if CUDAPP_CUDA_VERSION >= 4000 - .def("_launch_kernel", &cl::launch_kernel) + .def("_launch_kernel", &cl::launch_kernel, + (py::arg("grid_dim"), py::arg("block_dim"), + py::arg("parameter_buffer"), py::arg("shared_mem_bytes"), + py::arg("stream"), py::arg("cooperative")=false)) #endif ; } diff --git a/test/test_driver.py b/test/test_driver.py index fc689a02..66543b83 100644 --- a/test/test_driver.py +++ b/test/test_driver.py @@ -856,6 +856,87 @@ def test_prepared_invocation(self): drv.memcpy_dtoh(a_quadrupled, a_gpu) assert la.norm(a_quadrupled[1:] - 4 * a[1:]) == 0 + @mark_cuda_test + def test_cooperative_kernel(self): + if drv.get_version() < (9,): + pytest.skip("cooperative kernel launches require CUDA 9.0") + + dev = drv.Context.get_device() + if not dev.get_attribute(drv.device_attribute.COOPERATIVE_LAUNCH): + pytest.skip("device does not support cooperative kernel launches") + + source = r""" + #include + + extern "C" __global__ void cooperative_test(int *result, int value) + { + cooperative_groups::grid_group grid = cooperative_groups::this_grid(); + + if (grid.thread_rank() == 0) + { + result[0] = grid.is_valid(); + result[1] = value; + } + grid.sync(); + if (grid.thread_rank() == grid.size()-1) + result[2] = result[1]+1; + } + """ + + if drv.get_version() < (11,): + from pycuda.compiler import DynamicSourceModule + + mod = DynamicSourceModule(source, no_extern_c=True) + else: + mod = SourceModule(source, no_extern_c=True) + + func = mod.get_function("cooperative_test") + block = (32, 1, 1) + grid = (2, 1, 1) + expected = np.array([1, 41, 42], dtype=np.int32) + + result = np.zeros_like(expected) + func( + drv.Out(result), np.int32(41), block=block, grid=grid, cooperative=True + ) + assert (result == expected).all() + + func.prepare("Pi") + result_gpu = drv.mem_alloc(result.nbytes) + + drv.memset_d32(result_gpu, 0, result.size) + func.prepared_call(grid, block, result_gpu, 41, cooperative=True) + drv.memcpy_dtoh(result, result_gpu) + assert (result == expected).all() + + drv.memset_d32(result_gpu, 0, result.size) + get_time = func.prepared_timed_call( + grid, block, result_gpu, 41, cooperative=True + ) + assert get_time() >= 0 + drv.memcpy_dtoh(result, result_gpu) + assert (result == expected).all() + + stream = drv.Stream() + drv.memset_d32_async(result_gpu, 0, result.size, stream) + func.prepared_async_call( + grid, block, stream, result_gpu, 41, cooperative=True + ) + stream.synchronize() + drv.memcpy_dtoh(result, result_gpu) + assert (result == expected).all() + + blocks_per_mp = func.get_max_active_blocks_per_multiprocessor(block[0]) + mp_count = dev.get_attribute(drv.device_attribute.MULTIPROCESSOR_COUNT) + with pytest.raises(drv.LaunchError): + func.prepared_call( + (blocks_per_mp * mp_count + 1, 1, 1), + block, + result_gpu, + 41, + cooperative=True, + ) + @mark_cuda_test def test_prepared_with_vector(self): cuda_source = r"""