141 lines
5.7 KiB
Python
141 lines
5.7 KiB
Python
#################################################################################################
|
|
#
|
|
# Copyright (c) 2017 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
#
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
#
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
#
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
#
|
|
# 3. Neither the name of the copyright holder nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
#
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
#
|
|
#################################################################################################
|
|
|
|
import ctypes
|
|
from cutlass_cppgen.utils.lazy_import import lazy_import
|
|
cuda = lazy_import("cuda.cuda")
|
|
|
|
from cutlass_cppgen.backend.utils.device import device_cc
|
|
|
|
_supports_cluster_launch = None
|
|
|
|
|
|
def supports_cluster_launch():
|
|
from cuda import __version__
|
|
_version_splits = [int(x) for x in __version__.split("rc")[0].split(".post")[0].split(".")]
|
|
global _supports_cluster_launch
|
|
if _supports_cluster_launch is None:
|
|
major, minor = _version_splits[0], _version_splits[1]
|
|
_supports_cluster_launch = device_cc() in [90, 100, 101, 103] and (major > 11 or (major == 11 and minor >= 8))
|
|
return _supports_cluster_launch
|
|
|
|
|
|
class LaunchConfiguration:
|
|
def __init__(self, grid=[1, 1, 1], block=[1, 1, 1], smem=0):
|
|
self.grid = grid
|
|
self.block = block
|
|
self.shared_memory_capacity = smem
|
|
|
|
|
|
class ExecutableOperation:
|
|
def __init__(self, operation):
|
|
self.operation = operation
|
|
self.module = None
|
|
self.kernel = None
|
|
|
|
def name(self):
|
|
return self.operation.procedural_name()
|
|
|
|
def emit(self):
|
|
return ""
|
|
|
|
def can_implement(self, configuration, arguments):
|
|
raise NotImplementedError()
|
|
|
|
def get_host_workspace_size(self, arguments):
|
|
raise NotImplementedError()
|
|
|
|
def get_device_workspace_size(self, arguments):
|
|
raise NotImplementedError()
|
|
|
|
def plan(self, arguments):
|
|
raise NotImplementedError()
|
|
|
|
def initialize(self, host_workspace, device_workspace, launch_config, arguments, stream=None):
|
|
raise NotImplementedError()
|
|
|
|
def run_with_clusters(self, launch_config, kernel_params, stream=None):
|
|
if not stream:
|
|
stream = cuda.CUstream(0)
|
|
if hasattr(self.operation, "tile_description") and hasattr(self.operation.tile_description, "cluster_shape"):
|
|
attr = cuda.CUlaunchAttribute()
|
|
attr.value.clusterDim.x, attr.value.clusterDim.y, attr.value.clusterDim.z = self.operation.tile_description.cluster_shape
|
|
attr.id = cuda.CUstreamAttrID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
|
|
attrs = [attr]
|
|
|
|
# Allow for non-portable cluster sizes
|
|
err, = cuda.cuFuncSetAttribute(
|
|
self.kernel, cuda.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1)
|
|
if err != cuda.CUresult.CUDA_SUCCESS:
|
|
return err
|
|
else:
|
|
attrs = []
|
|
|
|
config = cuda.CUlaunchConfig()
|
|
config.gridDimX, config.gridDimY, config.gridDimZ = launch_config.grid
|
|
config.blockDimX, config.blockDimY, config.blockDimZ = launch_config.block
|
|
config.blockDimZ = launch_config.block[2]
|
|
config.sharedMemBytes = launch_config.shared_memory_capacity
|
|
config.hStream = stream
|
|
config.attrs = attrs
|
|
config.numAttrs = len(attrs)
|
|
|
|
err, = cuda.cuLaunchKernelEx(
|
|
config, f=self.kernel, kernelParams=kernel_params, extra=0)
|
|
return err
|
|
|
|
def run_without_clusters(self, launch_config, kernel_params, stream=None):
|
|
if not stream:
|
|
stream = cuda.CUstream(0)
|
|
err, = cuda.cuLaunchKernel(
|
|
self.kernel,
|
|
launch_config.grid[0], launch_config.grid[1], launch_config.grid[2],
|
|
launch_config.block[0], launch_config.block[1], launch_config.block[2],
|
|
launch_config.shared_memory_capacity,
|
|
stream,
|
|
kernel_params,
|
|
0)
|
|
|
|
return err
|
|
|
|
def run(self, host_workspace, device_workspace, launch_config, stream=None):
|
|
if not stream:
|
|
stream = cuda.CUstream(0)
|
|
cArg = (ctypes.c_char * len(host_workspace)).from_buffer(host_workspace)
|
|
packed = (ctypes.c_void_p * 1)()
|
|
packed[0] = ctypes.addressof(cArg)
|
|
|
|
if supports_cluster_launch():
|
|
return self.run_with_clusters(launch_config, packed, stream)
|
|
else:
|
|
return self.run_without_clusters(launch_config, packed, stream)
|