471 lines
14 KiB
Python
471 lines
14 KiB
Python
# 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()]
|