leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Release v4.0.0 (#2294)

Browse Source
This commit is contained in:
Kihiro Bando
2025-05-13 15:55:29 -04:00
committed by GitHub
parent ad7b2f5e84
commit f115c3f854
299 changed files with 51495 additions and 4413 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
python/CuTeDSL/base_dsl/runtime/__init__.py Normal file
View File

@@ -0,0 +1,29 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
"""
This module provides a runtime utility functions that are needed for
the DSL.
"""
from . import device_tensor
from . import dlpack_types
from . import cuda
from . import tensor_descriptor
from . import jit_arg_adapters
__all__ = [
"device_tensor",
"dlpack_types",
"cuda",
"tensor_descriptor",
"jit_arg_adapters",
]

470
python/CuTeDSL/base_dsl/runtime/cuda.py Normal file
View File

@@ -0,0 +1,470 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
"""
This module provides CUDA Python helper functions
"""
from functools import lru_cache
from dataclasses import dataclass
from typing import List, Optional
import numpy as np
import os
import ctypes
import cuda.bindings.driver as cuda
import cuda.bindings.nvrtc as nvrtc
# MLIR imports
from ..._mlir import ir
from ..._mlir.dialects import gpu
# Local module imports
from ..utils.logger import log as _log
from ..common import *
from .jit_arg_adapters import JitArgAdapterRegistry
# =============================================================================
# Utils
# =============================================================================
def _cudaGetErrorEnum(error):
if isinstance(error, cuda.CUresult):
err, name = cuda.cuGetErrorName(error)
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
elif isinstance(error, nvrtc.nvrtcResult):
return nvrtc.nvrtcGetErrorString(error)[1]
else:
raise DSLRuntimeError("Unknown error type: {}".format(error))
def _get_gpu_arch_info(major, minor):
"""Get GPU architecture information and compatibility details."""
gpu_arch_map = {
(7, 0): ("Volta", "sm_70", ["sm_70"]), # V100
(7, 5): ("Turing", "sm_75", ["sm_75"]), # RTX 20 Series, Quadro RTX
(8, 0): ("Ampere", "sm_80", ["sm_80"]), # A100
(8, 6): ("Ampere", "sm_86", ["sm_86", "sm_80"]), # RTX 30 Series
(8, 9): ("Ada", "sm_89", ["sm_89", "sm_86"]), # RTX 40 Series
(8, 7): ("Ampere", "sm_87", ["sm_87", "sm_86", "sm_80"]), # A10, A40
(9, 0): ("Hopper", "sm_90a", ["sm_90a"]), # H100
(10, 0): ("Blackwell", "sm_100a", ["sm_100a"]), # B200
}
return gpu_arch_map.get(
(major, minor), ("Unknown", f"sm_{major}{minor}", [f"sm_{major}{minor}"])
)
def get_compute_capability_major_minor(device_id: int = 0):
"""
Returns the compute capability of the CUDA device as a tuple of (major, minor).
For example: (8, 0) for Ampere, (9, 0) for Hopper, (10, 0) for Blackwell.
Returns None on failure.
"""
try:
checkCudaErrors(cuda.cuInit(0))
device = checkCudaErrors(cuda.cuDeviceGet(device_id))
major = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR,
device,
)
)
minor = checkCudaErrors(
cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
device,
)
)
return major, minor
except RuntimeError as e:
_log().info(f"Failed to get CUDA compute capability: {e}")
return None, None
@dataclass
class DeviceInfo:
"""Data class to store CUDA device information."""
device_count: int = 0
current_device: int = 0
device_name: Optional[str] = None
major_version: Optional[int] = None
minor_version: Optional[int] = None
arch_name: Optional[str] = None
sm_arch: Optional[str] = None
compatible_archs: Optional[List[str]] = None
memory_gb: Optional[float] = None
target_arch: Optional[str] = None
error_message: Optional[str] = None
initialization_failed: bool = False
def pretty_str(self) -> str:
"""
Convert DeviceInfo to a formatted string for display.
"""
info = ""
if self.initialization_failed:
return f"{Colors.BOLD}- CUDA initialization failed{Colors.RESET}"
if self.error_message:
return f"{Colors.BOLD}- Failed to get GPU info: {self.error_message}{Colors.RESET}"
if self.device_count > 0:
info += f"{Colors.BOLD}- CUDA devices available: {self.device_count} (current: {self.current_device})\n"
if self.major_version is not None and self.minor_version is not None:
info += f"- Architecture: {Colors.BLUE}{self.arch_name}{Colors.RESET} ({Colors.GREEN}{self.sm_arch}{Colors.RESET})\n"
info += f"- Compatible SM archs: {Colors.GREEN}{', '.join(self.compatible_archs or [])}{Colors.RESET}\n"
if self.memory_gb is not None:
info += f"- Total Memory: {Colors.BLUE}{self.memory_gb:.2f} GB{Colors.RESET}\n"
else:
info += f"- Compute capability: unknown\n"
info += f"- SM arch: unknown{Colors.RESET}\n"
else:
info += f"- No devices available\n"
return info
def get_device_info() -> DeviceInfo:
"""
Get detailed information about CUDA devices.
Returns a DeviceInfo dataclass with device information.
"""
device_info = DeviceInfo()
# Initialize CUDA if not already initialized
try:
result = cuda.cuInit(0)
if result[0].value: # Check for error
device_info.initialization_failed = True
return device_info
except:
pass
try:
# Get device count
result = cuda.cuDeviceGetCount()
device_info.device_count = result[1] if result[0].value == 0 else 0
if device_info.device_count > 0:
# Get current device
try:
result = cuda.cuCtxGetDevice()
if result[0].value == 0:
device_info.current_device = result[1]
except:
pass
# Get device name
try:
name_result = cuda.cuDeviceGetName(100, device_info.current_device)
if name_result[0].value == 0:
device_info.device_name = name_result[1]
except:
pass
# Get compute capability and architecture info
try:
major, minor = get_compute_capability_major_minor(
device_info.current_device
)
# Check if we successfully got the compute capability
if major is not None and minor is not None:
device_info.major_version = major
device_info.minor_version = minor
arch_name, sm_arch, compatible_archs = _get_gpu_arch_info(
device_info.major_version, device_info.minor_version
)
device_info.arch_name = arch_name
device_info.sm_arch = sm_arch
device_info.compatible_archs = compatible_archs
# Get memory info
try:
total_mem = cuda.cuDeviceGetAttribute(
cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_TOTAL_MEMORY,
device_info.current_device,
)
if total_mem[0].value == 0:
device_info.memory_gb = total_mem[1] / (
1024 * 1024 * 1024
) # Convert to GB
except:
pass
except Exception as e:
pass # Compute capability info will remain None
except Exception as e:
device_info.error_message = str(e)
return device_info
def checkCudaErrors(result):
"""Check CUDA errors and provide detailed error messages."""
if result[0].value:
error_code = result[0].value
error_name = _cudaGetErrorEnum(result[0])
raise DSLCudaRuntimeError(error_code, error_name)
if len(result) == 1:
return None
elif len(result) == 2:
return result[1]
else:
return result[1:]
# =============================================================================
# Driver Helpers
# =============================================================================
@lru_cache(maxsize=1)
def initialize_cuda_context(device_id: int = 0, flags: int = 0):
"""
Initializes the CUDA context for a specified device.
"""
# Initialize CUDA Driver API
_log().info(f"cuInit {flags}")
checkCudaErrors(cuda.cuInit(flags))
# Retrieve handle for device
_log().info(f"cuDeviceGet {device_id}")
cuDevice = checkCudaErrors(cuda.cuDeviceGet(device_id))
_log().info(f"{cuDevice} <-- cuDeviceGet")
# Create context
_log().info(f"cuCtxCreate {0} {cuDevice}")
context = checkCudaErrors(cuda.cuCtxCreate(0, cuDevice))
_log().info(f"{context} <-- cuCtxCreate")
return context
def load_cubin_module(cubin_file):
"""
Loads a CUBIN file and returns the module.
"""
# Load CUBIN file as binary data
_log().info(f"read cubin {cubin_file}")
with open(cubin_file, "rb") as f:
cubin_data = f.read()
# Load module data
_log().info(f"cuModuleLoadData {np.char.array(cubin_data).ctypes.data}")
module = checkCudaErrors(
cuda.cuModuleLoadData(np.char.array(cubin_data).ctypes.data)
)
return module
def unload_cubin_module(module):
"""
Unloads a CUBIN module.
"""
_log().info(f"cuModuleUnload {module}")
checkCudaErrors(cuda.cuModuleUnload(module))
def load_cubin_module_data(cubin_data):
"""
Loads a CUBIN from data and returns the module.
"""
# Load module data
_log().info(f"cuModuleLoadData {np.char.array(cubin_data).ctypes.data}")
module = checkCudaErrors(
cuda.cuModuleLoadData(np.char.array(cubin_data).ctypes.data)
)
return module
def get_kernel_function(module, kernel_name):
"""
Retrieves the kernel function from the module.
"""
_log().info(f"cuModuleGetFunction {module} {kernel_name}")
kernel = checkCudaErrors(
cuda.cuModuleGetFunction(module, bytes(kernel_name, "utf-8"))
)
_log().info(f"{kernel} <-- cuModuleGetFunction")
return kernel
def launch_kernel(kernel, grid_dims, block_dims, stream, smem_size=0, kernel_args=None):
"""
Launches the CUDA kernel.
"""
_log().info(
f"cuLaunchKernel {kernel} grid={grid_dims} blocks={block_dims} smem_size={smem_size} stream={stream} {kernel_args}"
)
checkCudaErrors(
cuda.cuLaunchKernel(
kernel,
grid_dims[0],
grid_dims[1],
grid_dims[2],
block_dims[0],
block_dims[1],
block_dims[2],
smem_size, # Shared memory size
stream,
kernel_args,
0, # Extra parameters
)
)
def stream_sync(stream):
"""
Synchronizes the CUDA stream.
"""
_log().info(f"cuStreamSynchronize {stream}")
checkCudaErrors(cuda.cuStreamSynchronize(stream))
def stream_create(id=0):
"""
Creates the CUDA stream.
"""
_log().info(f"cuStreamCreate {id}")
stream = checkCudaErrors(cuda.cuStreamCreate(id))
_log().info(f"{stream} <-- cuStreamCreate")
return stream
def stream_destroy(stream):
"""
Destroys the CUDA stream.
"""
_log().info(f"cuStreamDestroy {stream}")
checkCudaErrors(cuda.cuStreamDestroy(stream))
def context_destroy(context):
"""
Destroys the CUDA context.
"""
_log().info(f"cuCtxDestroy {context}")
checkCudaErrors(cuda.cuCtxDestroy(context))
def allocate(size_in_bytes: int, stream=None):
"""
Allocate device memory based on numpy host array size.
"""
_log().info("Allocate size_in_bytes=[%s] stream=[%s]", size_in_bytes, stream)
if stream is None:
device_memory = checkCudaErrors(cuda.cuMemAlloc(size_in_bytes))
else:
device_memory = checkCudaErrors(cuda.cuMemAllocAsync(size_in_bytes, stream))
_log().info("Allocated [%s]", device_memory)
return device_memory
def deallocate(device_pointer, stream=None):
"""
Deallocate the specified device memory pointer.
"""
_log().info(
"Deallocate device_pointer=[%s] stream=[%s]", hex(int(device_pointer)), stream
)
if stream is None:
checkCudaErrors(cuda.cuMemFree(device_pointer))
else:
checkCudaErrors(cuda.cuMemFreeAsync(device_pointer, stream))
def memcpy_h2d(host_pointer, device_pointer, size_in_bytes, stream=None):
"""
Copy data from host to device memory.
"""
_log().info(
"Copy host-to-device host_pointer[%s] device_ptr=[%s] size_in_bytes=[%s] stream=[%s]",
hex(host_pointer),
hex(int(device_pointer)),
size_in_bytes,
stream,
)
if stream is None:
checkCudaErrors(cuda.cuMemcpyHtoD(device_pointer, host_pointer, size_in_bytes))
else:
checkCudaErrors(
cuda.cuMemcpyHtoDAsync(device_pointer, host_pointer, size_in_bytes, stream)
)
def memcpy_d2h(host_pointer, device_pointer, size_in_bytes, stream=None):
"""
Copy data from device to host memory.
"""
_log().info(
"Copy device-host-to device_pointer=[%s] host_pointer[%s] size_in_bytes=[%s] stream=[%s]",
hex(int(device_pointer)),
hex(host_pointer),
size_in_bytes,
stream,
)
if stream is None:
checkCudaErrors(cuda.cuMemcpyDtoH(host_pointer, device_pointer, size_in_bytes))
else:
checkCudaErrors(
cuda.cuMemcpyDtoHAsync(host_pointer, device_pointer, size_in_bytes, stream)
)
def default_stream():
return cuda.CUstream(0)
def get_driver_version():
"""
Returns the CUDA driver version.
"""
return checkCudaErrors(cuda.cuDriverGetVersion())
def set_kernel_attribute(kernel, attribute, value):
"""
Sets a CUDA kernel attribute.
"""
return checkCudaErrors(cuda.cuFuncSetAttribute(kernel, attribute, value))
@JitArgAdapterRegistry.register_jit_arg_adapter(cuda.CUstream)
class StreamAdapter:
"""
Convert a CUDA stream to a stream representation for JIT arg generation.
"""
def __init__(self, arg):
self._arg = arg
self._c_pointer = ctypes.cast(self._arg.getPtr(), ctypes.c_void_p)
def __new_from_mlir_values__(self, values):
assert len(values) == 1
return values[0]
def __c_pointers__(self):
return [self._c_pointer]
def __get_mlir_types__(self):
return [gpu.AsyncTokenType.get()]

121
python/CuTeDSL/base_dsl/runtime/device_tensor.py Normal file
View File

@@ -0,0 +1,121 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import copy
from . import cuda as cuda_helpers
from .tensor_descriptor import *
from ..common import *
def allocate(tensor: TensorDescriptor, stream=None):
"""
Allocates GPU memory
"""
if tensor._check_is_managed_by_framework():
raise DSLRuntimeError(
"GPU tensors are managed by the framework and cannot be modified."
)
if not tensor.device_pointer is None:
raise DSLRuntimeError("Tensor is already allocated on the device.")
tensor.device_pointer = cuda_helpers.allocate(tensor.size_in_bytes, stream)
log().info("Allocate done tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer)
def deallocate(tensor: TensorDescriptor, stream=None):
"""
Deallocates GPU memory
"""
if tensor._check_is_managed_by_framework():
raise DSLRuntimeError(
"GPU tensors are managed by the framework and cannot be modified."
)
if tensor.device_pointer is None:
raise DSLRuntimeError("Tensor is not allocated on the device.")
log().info(
"Deallocating done tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer
)
cuda_helpers.deallocate(tensor.device_pointer, stream)
tensor.device_pointer = None
def copy_to_gpu(tensor: TensorDescriptor, do_allocate=True, stream=None):
"""
Copies data from host memory to the GPU memory.
If do_allocate is True, it first calls allocate
"""
log().info("copyin tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer)
if do_allocate:
allocate(tensor, stream)
cuda_helpers.memcpy_h2d(
tensor.data_ptr, tensor.device_pointer, tensor.size_in_bytes, stream
)
log().info("copyin done tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer)
return tensor
def copy_from_gpu(tensor: TensorDescriptor, do_deallocate=True, stream=None):
"""
Copies data from GPU memory back to the host.
If do_deallocate is True, it calls deallocate
"""
log().info("copyout tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer)
if tensor._check_is_managed_by_framework():
raise DSLRuntimeError(
"GPU tensors are managed by the framework and cannot be modified."
)
if tensor.device_pointer is None:
raise DSLRuntimeError("Tensor is not allocated on the device.")
cuda_helpers.memcpy_d2h(
tensor.data_ptr, tensor.device_pointer, tensor.size_in_bytes, stream
)
if do_deallocate:
deallocate(tensor, stream)
log().info("copyout done tensor=[%s] dev_ptr=[%s]", tensor, tensor.device_pointer)
def to_gpu(tensor, stream=None) -> TensorDescriptor:
"""
Copies the tensor to the GPU memory from Host memory
"""
if isinstance(tensor, TensorDescriptor):
new_tensor = copy.copy(tensor)
copy_to_gpu(new_tensor, stream=stream)
return new_tensor
if TensorDescriptor.can_transformed_to_dlpack(tensor):
new_tensor = TensorDescriptor(tensor)
copy_to_gpu(new_tensor, stream=stream)
return new_tensor
raise DSLRuntimeError("Unsupported type")
def from_gpu(tensor, stream=None) -> TensorDescriptor:
"""
Copies the tensor to the GPU memory from Host memory
"""
if isinstance(tensor, TensorDescriptor):
new_tensor = copy.copy(tensor)
copy_from_gpu(new_tensor, stream=stream)
return new_tensor
if TensorDescriptor.can_transformed_to_dlpack(tensor):
new_tensor = TensorDescriptor(tensor)
copy_from_gpu(new_tensor, stream=stream)
return new_tensor
raise DSLRuntimeError("Unsupported type")

76
python/CuTeDSL/base_dsl/runtime/dlpack_types.py Normal file
View File

@@ -0,0 +1,76 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
"""
This module provides helper structs for dlpack.
DLPack is an open standard for in-memory tensor structures, enabling
seamless sharing of tensors across different frameworks.
Learn more at: https://github.com/dmlc/dlpack
"""
import ctypes
import enum
class DLDeviceType(enum.IntEnum):
"""Enums for device types based on the DLPack specification."""
kDLCPU = 1
kDLGPU = 2
kDLCPUPinned = 3
class DLDataTypeCode:
"""Enums for data type codes based on the DLPack specification.
see https://github.com/dmlc/dlpack/blob/main/include/dlpack/dlpack.h
"""
kDLInt = 0
kDLUInt = 1
kDLFloat = 2
kDLOpaqueHandle = 3
kDLBfloat = 4
kDLComplex = 5
kDLBool = 6
class DLDevice(ctypes.Structure):
"""Structure representing the device information in DLPack."""
_fields_ = [
("device_type", ctypes.c_int), # kDLCPU, kDLGPU, etc.
("device_id", ctypes.c_int), # Device ID (e.g., GPU ID)
]
class DLDataType(ctypes.Structure):
"""Structure representing the data type in DLPack."""
_fields_ = [
("code", ctypes.c_uint8), # Data type code (e.g., kDLFloat)
("bits", ctypes.c_uint8), # Number of bits per value
("lanes", ctypes.c_uint16), # Number of lanes
]
class DLTensor(ctypes.Structure):
"""Structure representing the DLTensor in DLPack."""
_fields_ = [
("data", ctypes.c_void_p), # Pointer to tensor data
("device", DLDevice), # Device info
("ndim", ctypes.c_int), # Number of dimensions
("dtype", DLDataType), # Data type
("shape", ctypes.POINTER(ctypes.c_int64)), # Shape of tensor
("strides", ctypes.POINTER(ctypes.c_int64)), # Strides of tensor
("byte_offset", ctypes.c_uint64), # Byte offset to tensor data
]

188
python/CuTeDSL/base_dsl/runtime/jit_arg_adapters.py Normal file
View File

@@ -0,0 +1,188 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
"""
This module provides runtime utilities for JIT argument conversion in DSL.
"""
from functools import wraps
from typing import get_origin
# Local modules imports
from ..common import DSLRuntimeError
from ..typing import (
Constexpr,
Int32,
Float32,
Boolean,
)
def is_arg_spec_constexpr(arg_spec, arg_name, arg_index, owning_func):
"""
Check if the argument spec is a constexpr.
"""
def _is_reserved_python_func_arg(arg_index, arg_name, func):
"""
Check if the argument is a reserved python function argument.
"""
if arg_index != 0:
return False
if arg_name == "self":
return True
is_classmethod = isinstance(func, classmethod) or (
hasattr(func, "__func__") and isinstance(func.__func__, classmethod)
)
return arg_name == "cls" and is_classmethod
return (
_is_reserved_python_func_arg(arg_index, arg_name, owning_func)
or (isinstance(arg_spec, type) and issubclass(arg_spec, Constexpr))
or (get_origin(arg_spec) is Constexpr)
)
def is_argument_constexpr(arg, arg_spec, arg_name, arg_index, owning_func):
"""
Check if the argument is a constexpr.
"""
def _is_type_argument(arg, arg_annotation):
"""
Check if the argument is a type argument like Type[X]
"""
return isinstance(arg, type) and (
arg_annotation is None or get_origin(arg_annotation) is type
)
return (
is_arg_spec_constexpr(arg_spec, arg_name, arg_index, owning_func)
or _is_type_argument(arg, arg_spec)
or arg is None
)
class JitArgAdapterRegistry:
"""
A registry to keep track of the JIT argument adapters.
An adapter is a callable that converts a Python type to a type with following protocols supported:
- JitArgument
- DynamicExpression
The converted type can then be further processed by DSL to generate arguments for JIT functions.
"""
# A dictionary with key=type and value=callable
jit_arg_adapter_registry = {}
@classmethod
def register_jit_arg_adapter(cls, *dargs, **dkwargs):
"""
Register a JIT argument adapter callable
This can be used as a decorator on any callable like:
@register_jit_arg_adapter(my_py_type)
def my_adapter_for_my_py_type(arg):
...
@register_jit_arg_adapter(my_py_type)
class MyAdapterForMyPythonType:
...
The adapters are registered per type. If a type is already registerd, an error will be raised.
"""
def decorator(*dargs, **dkwargs):
darg_python_ty = dargs[0]
@wraps(darg_python_ty)
def wrapper(*args, **kwargs):
if len(args) != 1 or not callable(args[0]):
raise DSLRuntimeError(
"a callable must be provided for registering JIT argument adapter"
)
adapter = args[0]
if darg_python_ty in cls.jit_arg_adapter_registry:
raise DSLRuntimeError(
f"JIT argument adapter for {darg_python_ty} is already registered!",
context={
"Registered adapter": cls.jit_arg_adapter_registry[
darg_python_ty
],
"Adapter to be registered": adapter,
},
)
cls.jit_arg_adapter_registry[darg_python_ty] = adapter
return adapter
return wrapper
if len(dargs) > 0:
return decorator(*dargs, **dkwargs)
else:
raise DSLRuntimeError(
"a Python type must be provided for registering JIT argument adapter"
)
@classmethod
def get_registered_adapter(cls, ty):
"""
Get the registered JIT argument adapter for the given type.
"""
return cls.jit_arg_adapter_registry.get(ty, None)
# =============================================================================
# JIT Argument Adapters
# =============================================================================
@JitArgAdapterRegistry.register_jit_arg_adapter(int)
@JitArgAdapterRegistry.register_jit_arg_adapter(float)
@JitArgAdapterRegistry.register_jit_arg_adapter(bool)
def _convert_python_scalar(arg):
"""
Convert a Python scalar to a DSL type.
"""
conversion_map = {
int: Int32,
float: Float32,
bool: Boolean,
}
return conversion_map.get(type(arg))(arg)
@JitArgAdapterRegistry.register_jit_arg_adapter(tuple)
@JitArgAdapterRegistry.register_jit_arg_adapter(list)
def _convert_python_sequence(arg):
"""
Go through each element in the sequence and convert it to a type that can be
further processed by DSL to generate the corresponding JIT argument(s).
"""
adapted_arg = []
for elem in arg:
adapter = JitArgAdapterRegistry.get_registered_adapter(type(elem))
if adapter is not None:
converted_elem = adapter(elem)
adapted_arg.append(converted_elem)
else:
# If no registered adapter is found, just return the original element
adapted_arg.append(elem)
assert len(adapted_arg) == len(arg)
return type(arg)(adapted_arg)

201
python/CuTeDSL/base_dsl/runtime/tensor_descriptor.py Normal file
View File

@@ -0,0 +1,201 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
# Helpers
import itertools, operator
import ctypes
from . import dlpack_types as _dpack
from .dlpack_runtime import (
dlpack_to_tensor_desc,
get_tensor_desc_data_ptr,
get_tensor_desc_is_in_device,
get_tensor_desc_element_type,
get_tensor_desc_shape,
get_tensor_desc_stride,
get_tensor_desc_element_size_in_bytes,
get_tensor_desc_ndim,
get_tensor_desc_dtype_code,
get_tensor_desc_dtype_bits,
get_tensor_desc_device_type,
get_tensor_desc_device_id,
)
from ..utils.logger import log
from ..common import *
from ..typing import (
Boolean,
Float8E5M2,
Int64,
Int32,
Int16,
Int8,
Uint64,
Uint32,
Uint16,
Uint8,
Float64,
Float32,
Float16,
BFloat16,
)
class TensorDescriptor:
def __init__(self, tensor):
"""Initialize with a tensor that supports the DLPack protocol.
Args:
tensor: Any tensor object that implements __dlpack__ and __dlpack_device__
"""
self.tensor = tensor
self._capsule = dlpack_to_tensor_desc(tensor)
self.data_ptr = get_tensor_desc_data_ptr(self._capsule)
self.device_type = get_tensor_desc_device_type(self._capsule)
self.device_type = _dpack.DLDeviceType(self.device_type)
if self.device_type == _dpack.DLDeviceType.kDLGPU:
self.device_pointer = self.data_ptr
elif self.device_type == _dpack.DLDeviceType.kDLCPU:
self.device_pointer = None
else:
raise DSLRuntimeError(
f"DLPack device type is not supported {self.dl_tensor.device.device_type}"
)
log().info("TensorDescriptor is created = [%s]", self)
@staticmethod
def can_transformed_to_dlpack(dl_tensor):
if not hasattr(dl_tensor, "__dlpack__") or not hasattr(
dl_tensor, "__dlpack_device__"
):
return False
return True
@property
def is_in_device(self):
"""Check if the tensor is stored on a device."""
return not self.device_pointer is None
@property
def device_id(self):
"""Return device id where tensor resides."""
if self.is_in_device:
return get_tensor_desc_device_id(self._capsule)
return -1
@property
def element_type(self):
"""Return the corresponding Python type based on DLPack dtype metadata."""
str_element_type = get_tensor_desc_element_type(self._capsule)
dtype_map = {
# bool is 8bit from numpy and torch
"Bool": Boolean,
"Int64": Int64,
"Int32": Int32,
"Int16": Int16,
"Int8": Int8,
"UInt64": Uint64,
"UInt32": Uint32,
"UInt16": Uint16,
"UInt8": Uint8,
"Float64": Float64,
"Float32": Float32,
"Float16": Float16,
"BFloat16": BFloat16,
"Float8E5M2": Float8E5M2,
}
if str_element_type not in dtype_map:
raise KeyError(
f"Unsupported element type in dlpack: '{str_element_type}'. Supported types are: {list(dtype_map.keys())}"
)
return dtype_map[str_element_type]
@property
def shape(self):
"""Return the shape of the tensor."""
return get_tensor_desc_shape(self._capsule)
@property
def rank(self):
"""Return the rank of the tensor."""
return get_tensor_desc_ndim(self._capsule)
@property
def strides(self):
"""Return the rank of the tensor."""
return get_tensor_desc_stride(self._capsule)
@property
def element_size_in_bytes(self):
"""Calculate the element size in bytes of the DLPack tensor."""
return get_tensor_desc_element_size_in_bytes(self._capsule)
@property
def size_in_bytes(self):
"""Calculate the total size in bytes of the DLPack tensor."""
# Calculate the number of elements using the shape
ndim = get_tensor_desc_ndim(self._capsule)
shape = get_tensor_desc_shape(self._capsule)
num_elements = 1
for i in range(ndim):
num_elements *= shape[i]
# Total bytes
total_bytes = self.element_size_in_bytes * num_elements
return total_bytes
def __str__(self):
"""Return a compact string representation of the device_tensor with a tensor prefix."""
# Extract shape
shape = "x".join(map(str, self.shape))
# Extract dtype
dtype_code = get_tensor_desc_dtype_code(self._capsule)
dtype_bits = get_tensor_desc_dtype_bits(self._capsule)
dtype = (
f"i{dtype_bits}"
if dtype_code == _dpack.DLDataTypeCode.kDLInt
else f"f{dtype_bits}"
)
# Extract device
device_type = "cpu" if not self.is_in_device else "gpu"
return f"tensor<{shape}x{dtype}>_{device_type}"
def _check_is_managed_by_framework(self):
"""
Ensure the tensor is not managed by the framework (e.g., GPU tensor).
Raises an exception if the tensor is framework-managed.
"""
return self.device_type == _dpack.DLDeviceType.kDLGPU
def from_tensor(tensor) -> TensorDescriptor:
"""Create a TensorDescriptor from a tensor object."""
return TensorDescriptor(tensor)
def to_tensor(tensor_descriptor: TensorDescriptor):
"""Return tensor object from tensor descriptor."""
return tensor_descriptor.tensor
def is_tensor_descriptor(maybe_tensor_descriptor) -> bool:
"""Check if the object is a TensorDescriptor."""
return isinstance(
maybe_tensor_descriptor, TensorDescriptor
) or TensorDescriptor.can_transformed_to_dlpack(maybe_tensor_descriptor)