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

v4.1 release

Browse Source
This commit is contained in:
Junkai-Wu
2025-07-03 20:07:53 +08:00
committed by GitHub
parent b995f93317
commit a1aaf2300a
155 changed files with 18407 additions and 6068 deletions
Download Patch File Download Diff File Expand all files Collapse all files

186
python/CuTeDSL/base_dsl/typing.py
View File

@@ -46,29 +46,75 @@ from .._mlir.dialects import arith, math
@runtime_checkable
class DynamicExpression(Protocol):
"""
This is a protocol class that provides a common interface
to generate user-defined dynamic expressions.
"""Protocol defining the interface for object holding dynamic values in the DSL.
The DSL checks this protocol to determine if a class is a dynamic expression (SSA value) or not.
This protocol enables classes to represent dynamic values in the DSL. Classes implementing
this protocol can be used in JIT-compiled functions and dynamic value generation.
It is required for custom data types to work correctly with following JIT features:
* as function argument to call another JIT function from JIT function
* as return value from JIT function
* for constructions like if-else, while-loop, etc.
:param value: The MLIR operation result value to initialize the object with
:type value: ir.Value
**Required Methods**
* ``__extract_mlir_values__``: Extract MLIR values from the object
* ``__new_from_mlir_values__``: Create new instance from MLIR values
**Implementation Example**
To implement a custom data type that works with the DSL:
.. code-block:: python
class CustomData(metaclass=DslType):
def __init__(self, int_value):
self.int_value = int_value
def __extract_mlir_values__(self):
return [self.int_value]
def __new_from_mlir_values__(self, values):
return CustomData(values[0])
**Usage in JIT Functions**
When used in JIT-compiled functions, the DSL automatically extracts MLIR values:
.. code-block:: python
@jit
def caller():
x = CustomData(1)
return foo(x)
This generates MLIR like:
.. code-block:: mlir
func @caller() -> i32 {
%0 = func.call @foo(%arg0) : (i32) -> i32
return %0 : i32
}
"""
def __extract_mlir_values__(self):
"""
Generate a dynamic expression for the current object.
"""Extract MLIR values from this object.
:return: List of MLIR values
:return: List of MLIR values representing this object's data
:rtype: List[ir.Value]
"""
raise NotImplementedError
def __new_from_mlir_values__(self, values):
"""
Create a new object from MLIR values.
"""Create a new instance from MLIR values.
:param values: List of MLIR values
:param values: List of MLIR values to construct the object from
:type values: List[ir.Value]
:return: A new instance of the class that implements this protocol
:return: New instance of the implementing class
:rtype: Any
"""
raise NotImplementedError
@@ -77,50 +123,73 @@ class DynamicExpression(Protocol):
@runtime_checkable
class JitArgument(Protocol):
"""
This is a protocol class that provides a common interface
for JIT function arguments generation for Python to call JIT functions.
Protocol class defining the interface for JIT function argument generation.
The DSL checks this protocol to determine if a class is capable of providing information
needed for generating JIT function arguments.
This protocol enables classes to provide the necessary information for generating
JIT function arguments and allow the DSL JIT executor to call JIT compiled functions.
See breakdowns below for JitArgument protocol based JIT function calls.
**Required Methods**
* ``__c_pointers__``: Returns ctypes pointers for runtime execution
* ``__get_mlir_types__``: Returns MLIR types for function definition
* ``__new_from_mlir_values__``: Creates new instances from MLIR values
**Example**
.. code-block:: python
class CustomData:
def __init__(self, int_value, ...):
self.int_value = int_value
...
def __c_pointers__(self):
return [ctypes.pointer(ctypes.c_int32(self.int_value)), ...]
def __get_mlir_types__(self):
return [ir.IntegerType.get(32), ...]
def __new_from_mlir_values__(self, values):
return CustomData(values[0], ...)
@jit
def foo(x: CustomData):
return x.int_value + 1
a = x.int_value + 1
...
# Emit: `%c0 = arith.constant(1, i32)`
c1 = const(1, Int32)
# `c1` tracks `%c0` defined outside of function body of `foo`
# `%c0` can't be used directly in function body of `foo`
x = CustomData(c1, ...)
# `CustomData` is an argument of `foo`
foo(CustomData(1, ...))
When called like ``y = foo(x)``, the following steps occur:
1. JIT compiler generates MLIR function definition using ``__get_mlir_types__``:
1. JIT compiler generates MLIR function definition using ``__get_mlir_types__``
.. code-block:: mlir
func @foo(%arg0: i32, ...) -> i32 {
func.func @foo(%arg0: i32, ...) {
...
return
}
2. Function is traced in Python, wrapping MLIR values with ``__new_from_mlir_values__``:
2. JIT function can't use values from Python, so it needs to reconstruct the object from
MLIR values, a.k.a `%arg0`, with ``__new_from_mlir_values__`` and pass it to `foo`.
Following code demonstrates how JIT compiler reconstructs the object and pass to Python.
.. code-block:: python
# Implementation of IR tracing
new_x = CustomData(ir.Value(%arg0), ...)
y = foo(new_x)
# `x.int_value` is %arg0 rather than `c1` defined outside
# `x.int_value` is %arg0 rather than `c1` defined by Python.
3. For Python runtime execution, JIT engine invokes compiled function using ``__c_pointers__``:
3. For Python runtime execution, JIT engine invokes compiled function using ``__c_pointers__``
pointing to the underlying data object passing to JIT compiled function.
.. code-block:: python
jit_engine.invoke(foo, concat([x.__c_pointers__(), ...]))
jit_engine.invoke(compiled_foo, concat([x.__c_pointers__(), ...]))
"""
def __c_pointers__(self):
@@ -224,47 +293,6 @@ class DslType(type):
:property mlir_type: Returns the corresponding MLIR type for this DSL type
:type mlir_type: Any
**Examples**
Define a custom data type:
.. code-block:: python
class CustomData(metaclass=DslType, ...):
def __init__(self, int_value, ...):
self.int_value = int_value
...
def __str__(cls):
return "CustomData[int, ...]"
def __c_pointers__(self):
return [ctypes.pointer(ctypes.c_int32(self.int_value)), ...]
def __get_mlir_types__(self):
return [_T.i32(), ...]
def __extract_mlir_values__(self):
return [self.int_value, ...]
def __new_from_mlir_values__(self, values):
return CustomData(values[0], ...)
For JIT function calls, MLIR values are extracted with ``__extract_mlir_values__``:
.. code-block:: python
@jit
def caller():
x = CustomData(1, ...)
return foo(x)
.. code-block:: mlir
func @caller() -> i32 {
%0 = func.call @foo(%arg0, ...) : (i32, ...) -> i32
return %0 : i32
}
"""
_is_abstract: bool
@@ -946,9 +974,12 @@ class Numeric(metaclass=NumericMeta, is_abstract=True):
:return: The result of the logical not operation
:rtype: Boolean
"""
ty = type(self)
zero_val = arith.constant(ty.mlir_type, ty.zero)
return self.__eq__(ty(zero_val), loc=loc, ip=ip)
if isinstance(self.value, (int, float, bool)):
return not self.value
else:
ty = type(self)
zero_val = arith.constant(ty.mlir_type, ty.zero)
return self.__eq__(ty(zero_val), loc=loc, ip=ip)
def __dsl_and__(self, other, *, loc=None, ip=None):
"""DSL implementation of Python's `and` operator.
@@ -1057,6 +1088,15 @@ class Numeric(metaclass=NumericMeta, is_abstract=True):
],
)
def __index__(self):
if isinstance(self.value, (int, float, bool)):
return self.value
else:
raise DSLRuntimeError(
f"'{type(self.value)}' object cannot be interpreted as an integer",
suggestion="Mark the loop as dynamic with `dynamic_expr` or `range_dynamic` and decorate the parent function with `jit` decorator",
)
def __neg__(self, *, loc=None, ip=None):
if isinstance(self, (bool, int, float)):
return type(self)(-self.value) # type: ignore
@@ -1813,7 +1853,7 @@ class IRVariadic:
def __init__(self, operands):
"""
Create a list of variadic operands. `operands` must be SSA values.
Create a list of variadic operands. `operands` must be dynamic values.
"""
self.operands = operands