432 lines
13 KiB
Python
432 lines
13 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.
|
|
|
|
"""Kernel specification classes for TVM-FFI function parameters."""
|
|
from abc import ABC
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Optional, Union
|
|
|
|
try:
|
|
import tvm_ffi
|
|
except ModuleNotFoundError:
|
|
pass
|
|
|
|
class DefaultConfig:
|
|
"""Default configuration with context manager support."""
|
|
|
|
_current: Optional["DefaultConfig"] = None
|
|
_old_current: Optional["DefaultConfig"] = None
|
|
device_type: str
|
|
|
|
def __init__(self, *, device_type: Optional[str] = None) -> None:
|
|
"""Initialize a default configuration.
|
|
|
|
Parameters
|
|
----------
|
|
device_type : str, optional
|
|
The device type (e.g., "cpu", "cuda", "metal"). If None, copies from current config.
|
|
"""
|
|
if device_type is None:
|
|
device_type = DefaultConfig.current().device_type # type: ignore[union-attr]
|
|
self.device_type = device_type
|
|
|
|
def __enter__(self) -> "DefaultConfig":
|
|
"""Enter the context manager."""
|
|
self._old_current = DefaultConfig._current
|
|
DefaultConfig._current = self
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: Optional[type],
|
|
exc_val: Optional[BaseException],
|
|
exc_tb: Optional[object],
|
|
) -> None:
|
|
"""Exit the context manager."""
|
|
DefaultConfig._current = self._old_current
|
|
|
|
@classmethod
|
|
def current(cls) -> Optional["DefaultConfig"]:
|
|
"""Get the current default configuration.
|
|
|
|
Returns
|
|
-------
|
|
Optional[DefaultConfig]
|
|
The current default configuration.
|
|
"""
|
|
return cls._current
|
|
|
|
@classmethod
|
|
def _set_init_default_config(cls) -> None:
|
|
"""Set the initial default configuration."""
|
|
current = cls.__new__(cls)
|
|
current.device_type = "cuda"
|
|
cls._current = current
|
|
|
|
|
|
# Initialize the default config to cuda
|
|
DefaultConfig._set_init_default_config()
|
|
|
|
|
|
class Param(ABC):
|
|
"""Base class for all parameters."""
|
|
|
|
|
|
class Var(Param):
|
|
"""variables: pointer, integer, floating-point, boolean, etc.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
dtype : str | tvm_ffi.dtype
|
|
The data type of the parameter.
|
|
divisibility: Optional[int]
|
|
The divisibility of the parameter, by default None.
|
|
|
|
"""
|
|
|
|
name: str
|
|
dtype: "tvm_ffi.dtype"
|
|
divisibility: Optional[int]
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
dtype: Union[str, "tvm_ffi.dtype"],
|
|
*,
|
|
divisibility: Optional[int] = None,
|
|
) -> None:
|
|
"""Initialize a Var parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
dtype : str | tvm_ffi.dtype
|
|
The data type of the parameter.
|
|
|
|
"""
|
|
self.name = name
|
|
self.dtype = tvm_ffi.dtype(dtype)
|
|
self.divisibility = divisibility
|
|
|
|
|
|
class Shape(Param):
|
|
"""Shape parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
shape: list[int | Var]
|
|
The shape of the parameter.
|
|
|
|
"""
|
|
|
|
name: str
|
|
shape: list[Union[int, Var]]
|
|
|
|
def __init__(self, name: str, shape: list[Union[int, Var]]) -> None:
|
|
"""Initialize a Shape parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
shape : list[int | Var]
|
|
The shape of the parameter.
|
|
|
|
"""
|
|
self.name = name
|
|
self.shape = shape
|
|
|
|
|
|
class Tensor(Param):
|
|
"""Tensor parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
dtype : str | tvm_ffi.dtype
|
|
The data type of the parameter.
|
|
shape : Sequence[int | Var]
|
|
The shape of the parameter.
|
|
device_type : int
|
|
The device type of the parameter.
|
|
device_id : Var
|
|
The device id of the parameter.
|
|
strides : Optional[Sequence[Var]], optional
|
|
The strides of the parameter, by default None.
|
|
"""
|
|
|
|
name: str
|
|
shape: list[Union[int, Var]]
|
|
dtype: "tvm_ffi.dtype"
|
|
strides: Optional[list[Var]]
|
|
dlpack_device_type: int
|
|
device_id: Var
|
|
map_tensor_dtype_f4x2_to_f4: bool
|
|
data_alignment: Optional[int]
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
shape: Sequence[Union[int, Var]],
|
|
dtype: Union[str, "tvm_ffi.dtype"],
|
|
*,
|
|
device_type: Optional[str] = None,
|
|
strides: Optional[Sequence[Var]] = None,
|
|
map_tensor_dtype_f4x2_to_f4: bool = False,
|
|
data_alignment: Optional[int] = None,
|
|
) -> None:
|
|
"""Initialize a Tensor parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
device_type : str
|
|
The device type of the parameter.
|
|
shape : Sequence[int | Var]
|
|
The shape of the parameter.
|
|
dtype : str | tvm_ffi.dtype
|
|
The data type of the parameter.
|
|
strides : Optional[Sequence[Var]], optional
|
|
The strides of the parameter, by default None.
|
|
map_tensor_dtype_f4x2_to_f4: bool
|
|
Whether to map tensor dtype float4x2 to float4 for internal use.
|
|
data_alignment: Optional[int], optional
|
|
The data alignment of the parameter, by default None.
|
|
"""
|
|
self.name = name
|
|
self.data = Var(name + ".data", tvm_ffi.dtype("handle"))
|
|
self.shape: list[Union[int, Var]] = list(shape)
|
|
self.dtype = tvm_ffi.dtype(dtype)
|
|
self.strides: Optional[list[Var]] = list(strides) if strides is not None else None
|
|
self.data_alignment = data_alignment
|
|
|
|
# Use default device type if none specified
|
|
if device_type is None:
|
|
device_type = DefaultConfig.current().device_type # type: ignore[union-attr]
|
|
|
|
example_device = tvm_ffi.device(device_type, 0)
|
|
self.dlpack_device_type = example_device.dlpack_device_type()
|
|
self.device_type_name = example_device.type
|
|
self.device_id = Var(name + ".device_id", tvm_ffi.dtype("int32"))
|
|
self.map_tensor_dtype_f4x2_to_f4 = map_tensor_dtype_f4x2_to_f4
|
|
|
|
|
|
class Stream(Param):
|
|
"""Stream parameter."""
|
|
|
|
name: str
|
|
var: Var
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Initialize a Stream parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
"""
|
|
self.name = name
|
|
self.var = Var(name, tvm_ffi.dtype("handle"))
|
|
|
|
|
|
class EnvStream(Param):
|
|
"""EnvStream parameter.
|
|
|
|
Note
|
|
----
|
|
This parameter spec indicates that we expect
|
|
to call the function with TVMFFIEnvGetStream to get the stream
|
|
and it is not part of the FFI function signature.
|
|
"""
|
|
|
|
name: str
|
|
var: Var
|
|
|
|
def __init__(self, name: str) -> None:
|
|
"""Initialize a EnvStream parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
"""
|
|
self.name = name
|
|
self.var = Var(name, tvm_ffi.dtype("handle"))
|
|
|
|
|
|
class DataPointer(Param):
|
|
"""Data pointer parameter.
|
|
|
|
The main difference between DataPointer and Var with handle type
|
|
is that DataPointer can contain address space information.
|
|
We also allow passing in int as data pointer, so it can
|
|
conveniently support torch.Tensor.data_ptr() as input.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
"""
|
|
|
|
name: str
|
|
var: Var
|
|
address_space: Optional[int]
|
|
|
|
def __init__(self, name: str, address_space: Optional[int] = None) -> None:
|
|
"""Initialize a DataPointer parameter.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The parameter name.
|
|
"""
|
|
self.name = name
|
|
self.var = Var(name, tvm_ffi.dtype("handle"))
|
|
self.address_space = address_space
|
|
|
|
|
|
def signature(name: str, params: list[Param]) -> str:
|
|
"""Generate a function signature string from name and parameters.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
The function name.
|
|
params : list[Param]
|
|
List of parameter objects (Var or Tensor).
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
The formatted function signature.
|
|
|
|
Raises
|
|
------
|
|
ValueError
|
|
If an unknown parameter type is encountered.
|
|
|
|
"""
|
|
param_strs = []
|
|
|
|
for param in params:
|
|
if isinstance(param, Var):
|
|
param_str = f"{param.name}: {param.dtype}"
|
|
elif isinstance(param, Tensor):
|
|
# Format tensor shape
|
|
shape_strs = []
|
|
for dim in param.shape:
|
|
if isinstance(dim, Var):
|
|
shape_strs.append(dim.name)
|
|
else:
|
|
shape_strs.append(str(dim))
|
|
shape_str = "[" + ", ".join(shape_strs) + "]"
|
|
param_str = f"{param.name}: Tensor({shape_str}, {param.dtype})"
|
|
elif isinstance(param, Shape):
|
|
# Format shape parameter
|
|
shape_strs = []
|
|
for dim in param.shape:
|
|
if isinstance(dim, Var):
|
|
shape_strs.append(dim.name)
|
|
else:
|
|
shape_strs.append(str(dim))
|
|
shape_str = "[" + ", ".join(shape_strs) + "]"
|
|
param_str = f"{param.name}: Shape({shape_str})"
|
|
elif isinstance(param, Stream):
|
|
param_str = f"{param.name}: Stream"
|
|
elif isinstance(param, DataPointer):
|
|
param_str = f"{param.name}: DataPointer"
|
|
elif isinstance(param, EnvStream):
|
|
# env stream is not part of the FFI function signature
|
|
# continue to skip append
|
|
continue
|
|
else:
|
|
raise TypeError(f"Unsupported parameter type: {type(param)}")
|
|
|
|
param_strs.append(param_str)
|
|
|
|
return f"{name}({', '.join(param_strs)})"
|
|
|
|
|
|
def create_map_tensor_dtype_f4x2_to_f4_spec(f4_tensor_spec: Tensor) -> Tensor:
|
|
"""
|
|
Create a new Tensor spec that can be translated to f4 via f4x2->f4 conversion.
|
|
|
|
Parameters
|
|
----------
|
|
f4_tensor_spec : Tensor
|
|
The original Tensor spec that declares a f4 tensor.
|
|
|
|
Returns
|
|
-------
|
|
Tensor
|
|
The new Tensor spec that can be translated to f4 via f4x2->f4 conversion.
|
|
"""
|
|
if f4_tensor_spec.dtype != tvm_ffi.dtype("float4_e2m1fn"):
|
|
raise ValueError("f4_tensor_spec must be a float4 tensor")
|
|
|
|
def find_stride_one_index() -> int:
|
|
if f4_tensor_spec.strides is None:
|
|
return len(f4_tensor_spec.shape) - 1
|
|
for i, stride in enumerate(f4_tensor_spec.strides):
|
|
if isinstance(stride, int) and stride == 1:
|
|
return i
|
|
raise ValueError("Cannot find dimension with stride=1")
|
|
stride_one_index = find_stride_one_index()
|
|
|
|
def divisibility_divide_by_2(value: Var) -> Optional[int]:
|
|
if value.divisibility is not None:
|
|
if value.divisibility % 2 != 0:
|
|
raise ValueError(f"Dimension with stride=1 must be divisible by 2")
|
|
return value.divisibility // 2
|
|
return None
|
|
|
|
def map_shape(index: int, value: Union[int, Var]) -> Union[int, Var]:
|
|
if index == stride_one_index:
|
|
if isinstance(value, int):
|
|
if value % 2 != 0:
|
|
raise ValueError(f"Dimension {index} with stride=1 must be even")
|
|
return value // 2
|
|
# create a new var with the same name and dtype
|
|
return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value))
|
|
return value
|
|
|
|
def map_stride(index: int, value: Union[int, Var]) -> Union[int, Var]:
|
|
if index != stride_one_index:
|
|
if isinstance(value, int):
|
|
if value % 2 != 0:
|
|
raise ValueError(f"Dimension {index} with stride != 1 must be even")
|
|
return value // 2
|
|
# create a new var with the same name and dtype
|
|
return Var(value.name, value.dtype, divisibility=divisibility_divide_by_2(value))
|
|
return value
|
|
|
|
new_shape = [map_shape(i, x) for i, x in enumerate(f4_tensor_spec.shape)]
|
|
if f4_tensor_spec.strides is not None:
|
|
new_strides = [map_stride(i, x) for i, x in enumerate(f4_tensor_spec.strides)]
|
|
else:
|
|
new_strides = None
|
|
|
|
return Tensor(
|
|
f4_tensor_spec.name,
|
|
new_shape,
|
|
dtype=tvm_ffi.dtype("float4_e2m1fnx2"),
|
|
strides=new_strides,
|
|
map_tensor_dtype_f4x2_to_f4=True,
|
|
)
|