Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -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()]
|
||||
Reference in New Issue
Block a user