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

396
python/CuTeDSL/base_dsl/ast_helpers.py
View File

@@ -15,6 +15,8 @@ The preprocessor read through python's ast and changes the input code.
"""
from typing import Callable, Iterator, Optional, overload
from typing_extensions import deprecated
import warnings
from .utils.logger import log
from .common import *
@@ -30,13 +32,9 @@ class Executor:
set_functions: Assigns the functions for checking loop bounds and
conditional evaluation.
for_dynamic: Generates MLIR for OP
for_constexpr: Executes a for loop at JIT compile-time
for_execute: Decides whether to execute the loop at compile-time or generate MLIR for OP based on the provided bounds.
if_dynamic: Generates MLIR if OP
if_constexpr: Executes a if at JIT compile-time by python interpreter
if_execute: Decides whether to execute the if statement at compile-time or generate MLIR if OP based on the predicate.
for_execute: Generates MLIR for OP
while_execute: Generates MLIR while OP
if_execute: generate MLIR if OP
"""
def __init__(self):
@@ -44,6 +42,9 @@ class Executor:
self._loop_execute_range_dynamic = None
self._if_dynamic = None
self._while_dynamic = None
self._compare_executor = None
self._any_executor = None
self._all_executor = None
def set_functions(
self,
@@ -51,11 +52,17 @@ class Executor:
loop_execute_range_dynamic: Callable,
if_dynamic: Callable,
while_dynamic: Callable,
compare_executor: Callable,
any_executor: Callable = None,
all_executor: Callable = None,
):
self._is_dynamic_expression = is_dynamic_expression
self._loop_execute_range_dynamic = loop_execute_range_dynamic
self._if_dynamic = if_dynamic
self._while_dynamic = while_dynamic
self._compare_executor = compare_executor
self._any_executor = any_executor
self._all_executor = all_executor
@staticmethod
def convert_to_list(x):
@@ -83,31 +90,6 @@ class Executor:
return res[0]
return res
def for_dynamic(
self,
func: Callable,
start,
stop,
step,
used_args: list,
iter_args: list,
iter_arg_names: list,
unroll=bool,
unroll_full=int,
):
log().debug("start [%s] stop [%s] step [%s]", start, stop, step)
return self._loop_execute_range_dynamic(
func,
start,
stop,
step,
used_args,
iter_args,
iter_arg_names,
unroll,
unroll_full,
)
@staticmethod
def for_constexpr(
func: Callable,
@@ -143,44 +125,14 @@ class Executor:
iter_arg_names=[],
unroll=-1,
unroll_full=False,
is_range_constexpr=None,
pipelining=None,
):
assert (
self._loop_execute_range_dynamic and self._is_dynamic_expression
self._loop_execute_range_dynamic
), "Functions must be set before execution."
log().debug("start [%s] stop [%s] step [%s]", start, stop, step)
any_dynamic_expression = (
self._is_dynamic_expression(start)
or self._is_dynamic_expression(stop)
or self._is_dynamic_expression(step)
)
if is_range_constexpr is None:
if not any_dynamic_expression:
return self.for_constexpr(func, start, stop, step, used_args, iter_args)
else:
return self.for_dynamic(
func,
start,
stop,
step,
used_args,
iter_args,
iter_arg_names,
unroll,
unroll_full,
)
# Ensure bounds are compile-time constants for constexpr execution
if is_range_constexpr:
if any_dynamic_expression:
raise DSLRuntimeError(
"Loop bounds must be constexpr (compile-time constants)"
)
return self.for_constexpr(func, start, stop, step, used_args, iter_args)
# MLIR generation
return self.for_dynamic(
return self._loop_execute_range_dynamic(
func,
start,
stop,
@@ -190,40 +142,9 @@ class Executor:
iter_arg_names,
unroll,
unroll_full,
pipelining,
)
def if_dynamic(
self,
pred,
then_block: Callable,
else_block: Optional[Callable] = None,
used_args=[],
yield_args=[],
yield_arg_names=[],
):
return self._if_dynamic(
pred, then_block, else_block, used_args, yield_args, yield_arg_names
)
@staticmethod
def if_constexpr(
pred,
then_block: Callable,
else_block: Optional[Callable] = None,
used_args=[],
yield_args=[],
):
if pred:
log().debug(" running then block [%s]", yield_args)
res = then_block(*used_args, *yield_args)
log().debug("result [%s]", res)
return Executor.converge_ret_val(res)
elif else_block is not None:
log().debug("running else [%s]", yield_args)
res = else_block(*used_args, *yield_args)
log().debug("result [%s]", res)
return Executor.converge_ret_val(res)
def if_execute(
self,
pred,
@@ -232,94 +153,14 @@ class Executor:
used_args=[],
yield_args=[],
yield_arg_names=[],
if_constexpr=None,
):
assert (
self._if_dynamic and self._is_dynamic_expression
), "Functions must be set before execution."
is_if_constexpr = not self._is_dynamic_expression(pred)
if if_constexpr is None:
if is_if_constexpr:
return self.if_constexpr(
pred, then_block, else_block, used_args, yield_args
)
else:
return self.if_dynamic(
pred, then_block, else_block, used_args, yield_args, yield_arg_names
)
# Ensure bounds are compile-time constants for constexpr execution
if if_constexpr:
if not is_if_constexpr:
raise DSLRuntimeError(
"If predicate must be constexpr (compile-time constants)"
)
return self.if_constexpr(
pred, then_block, else_block, used_args, yield_args
)
assert self._if_dynamic, "Functions must be set before execution."
# MLIR generation
return self.if_dynamic(
return self._if_dynamic(
pred, then_block, else_block, used_args, yield_args, yield_arg_names
)
def while_dynamic(
self,
while_before_block: Callable,
while_after_block: Callable,
used_args=[],
yield_args=[],
yield_arg_names=[],
):
return self._while_dynamic(
while_before_block,
while_after_block,
used_args,
yield_args,
yield_arg_names,
)
@staticmethod
def while_constexpr(
while_before_block,
while_after_block,
used_args=[],
yield_args=[],
):
log().debug(
"while_constexpr begin %s", while_before_block.__qualname__
)
cond, loop_results = while_before_block(*used_args, *yield_args)
while cond:
loop_results = Executor.convert_to_list(loop_results)
log().debug(
"calling while_after [%s], [%s]",
used_args,
loop_results,
)
loop_results = while_after_block(*used_args, *loop_results)
log().debug(
"while after [%s]", loop_results
)
loop_results = Executor.convert_to_list(loop_results)
log().debug(
"calling while_before [%s], [%s]",
used_args,
loop_results,
)
cond, loop_results = while_before_block(*used_args, *loop_results)
log().debug(
"while_before cond, results [%s], [%s]",
cond,
loop_results,
)
log().debug(
"while_constexpr results %s", loop_results
)
return Executor.converge_ret_val(loop_results)
def while_execute(
self,
pred,
@@ -328,26 +169,11 @@ class Executor:
used_args=[],
yield_args=[],
yield_arg_names=[],
while_constexpr=None,
):
assert (
self._while_dynamic and self._is_dynamic_expression
), "Functions must be set before execution."
is_while_constexpr = not self._is_dynamic_expression(pred)
# Ensure bounds are compile-time constants for constexpr execution
if while_constexpr:
if not is_while_constexpr:
raise DSLRuntimeError(
"While predicate must be constexpr (compile-time constants)"
)
return self.while_constexpr(
while_before_block, while_after_block, used_args, yield_args
)
assert self._while_dynamic, "Functions must be set before execution."
# MLIR generation
return self.while_dynamic(
return self._while_dynamic(
while_before_block,
while_after_block,
used_args,
@@ -367,15 +193,16 @@ def loop_selector(
start,
stop,
step,
*,
used_args=[],
iter_args=[],
iter_arg_names=[],
unroll=-1,
unroll_full=False,
constexpr=None,
pipelining=None,
):
log().debug(
"start [%s] stop [%s] step [%s] used_args [%s] iter_args [%s] unroll [%s] unroll_full [%s] constexpr [%s]",
"start [%s] stop [%s] step [%s] used_args [%s] iter_args [%s] unroll [%s] unroll_full [%s] pipelining [%s]",
start,
stop,
step,
@@ -383,7 +210,7 @@ def loop_selector(
iter_args,
unroll,
unroll_full,
constexpr,
pipelining,
)
from .typing import Integer, Numeric
@@ -408,7 +235,7 @@ def loop_selector(
iter_arg_names,
unroll,
unroll_full,
constexpr,
pipelining,
)
return ir_loop
@@ -443,7 +270,6 @@ def while_executor(
used_args=[],
yield_args=[],
yield_arg_names=[],
constexpr=None,
):
return executor.while_execute(
pred,
@@ -452,7 +278,6 @@ def while_executor(
used_args,
yield_args,
yield_arg_names,
constexpr,
)
@@ -463,10 +288,9 @@ def if_executor(
used_args=[],
yield_args=[],
yield_arg_names=[],
constexpr=None,
):
return executor.if_execute(
pred, then_block, else_block, used_args, yield_args, yield_arg_names, constexpr
pred, then_block, else_block, used_args, yield_args, yield_arg_names
)
@@ -475,75 +299,70 @@ def if_executor(
# =============================================================================
class range_dynamic:
class range:
@overload
def __new__(cls, stop, unroll=0, unroll_full=False):
def __new__(cls, stop, unroll=0, unroll_full=False, pipelining=None):
pass
@overload
def __new__(cls, start, stop, step, unroll=0, unroll_full=False):
def __new__(cls, start, stop, step, unroll=0, unroll_full=False, pipelining=None):
pass
def __new__(cls, *args, **kwargs):
raise DSLRuntimeError("range_dynamic should be always preprocessed to IR")
class range_constexpr:
def __init__(self, *args):
if len(args) == 1:
self.start = 0
self.stop = args[0]
self.step = 1
elif len(args) == 2:
self.start, self.stop = args
self.step = 1
elif len(args) == 3:
self.start, self.stop, self.step = args
else:
raise DSLRuntimeError(
"range_constexpr supports up to 3 arguments (start, stop, step)"
)
# Ensure the arguments are compile-time constants (if required)
for arg_name, arg_value in [
("step", self.step),
("start", self.start),
("stop", self.stop),
]:
if executor._is_dynamic_expression(arg_value):
raise DSLRuntimeError(
f"`range_constexpr` requires `constexpr` (non-IR Values) for all arguments, "
f"but `{arg_name}` is not. If the arguments are dynamic, use `range`; the DSL "
f"will handle them during runtime. ",
suggestion="Use `range` instead of `range_constexpr`.",
)
raise DSLRuntimeError("dynamic range should be always preprocessed to IR")
def __iter__(self) -> Iterator[int]:
current = self.start
while current < self.stop:
yield current
current += self.step
raise DSLRuntimeError("dynamic range should be always preprocessed to IR")
@deprecated(
"range_dynamic is deprecated and will be removed in the future, please remove it."
)
def range_dynamic(*args, **kwargs):
raise DSLRuntimeError("range_dynamic should be always preprocessed to IR")
def range_constexpr(*args):
raise DSLRuntimeError("range_constexpr should be preprocessed by preprocessor.")
# =============================================================================
# If expressions
# =============================================================================
def const_expr(expression):
if executor._is_dynamic_expression(expression):
"""
This function is used to check if the expression is a python value.
If the expression is a python value, return the boolean value of the expression.
If the expression is a dynamic expression, raise an error.
"""
from .typing import Numeric
failed = False
if isinstance(expression, Numeric):
if isinstance(expression.value, (int, float, bool)):
return expression.value
else:
failed = True
elif executor._is_dynamic_expression(expression):
failed = True
if failed:
raise DSLRuntimeError(
f"The function `const_expr({expression})` received a dynamic expression (non compile-time constant).",
context={
"const_expr": "Accepts only constexpr (compile-time constant)",
"If your expression depends on dynamic values": "Avoid marking it as `const_expr()`",
"If the expression could be either dynamic or constexpr": "Omit explicit `const_expr()` marker; the DSL will infer the correct handling automatically",
"If your expression depends on dynamic values": "Remove `const_expr()`",
},
)
return expression
@deprecated(
"dynamic_expr is deprecated and will be removed in the future, please remove it."
)
def dynamic_expr(expression):
raise DSLRuntimeError("dynamic_expr should be always preprocessed to IR")
return expression
# =============================================================================
@@ -582,3 +401,86 @@ def bool_cast(value):
suggestion = "Please explicitly convert to boolean with expressions like comparision."
)
return bool(value)
def compare_executor(left, comparators, ops):
"""
Executes comparison operations with a left operand and a list of comparators.
Args:
left: The leftmost value in the comparison chain
comparators: A list of values to compare against
ops: A list of comparison operators to apply
Returns:
The result of the comparison chain
Raises:
AssertionError: If the executor function is not set before execution
"""
assert (
executor._compare_executor is not None
), "Function must be set before execution."
return executor._compare_executor(left, comparators, ops)
def any_executor(iterable):
"""Executes the 'any' operation on an iterable, handling both dynamic and static expressions.
:param iterable: An iterable to check if any elements evaluate to True
:type iterable: Iterable
:return: boolean of Python value or IR value
:rtype: bool or cutlass.Boolean
"""
if executor._any_executor and executor._is_dynamic_expression(iterable):
return executor._any_executor(iterable)
else:
return any(iterable)
def all_executor(iterable):
"""Executes the 'all' operation on an iterable, handling both dynamic and static expressions.
:param iterable: An iterable to check if all elements evaluate to True
:type iterable: Iterable
:return: boolean of Python value or IR value
:rtype: bool or cutlass.Boolean
"""
if executor._all_executor and executor._is_dynamic_expression(iterable):
return executor._all_executor(iterable)
else:
return all(iterable)
# =============================================================================
# Control flow checks
# =============================================================================
def range_value_check(*args):
"""
Ensure all `range_constexpr` bounds are compile-time constants (Python ints).
"""
try:
return tuple(arg.__index__() for arg in args)
except:
raise DSLRuntimeError(
"`range_constexpr` requires constexpr (compile-time constant) for all arguments.",
suggestion="Use `range` instead of `range_constexpr`.",
)
def range_perf_warning(filename, lineno, *args):
has_dynamic_expr = False
for arg in args:
if executor._is_dynamic_expression(arg):
has_dynamic_expr = True
break
if not has_dynamic_expr:
warnings.warn_explicit(
(
"The loop was previously unrolled in Python, but now it may not unroll in IR. This may cause performance regression."
"If you want to unroll the loop in Python, please use `range_constexpr` instead of `range`."
),
category=UserWarning,
filename=filename,
lineno=lineno,
)

805
python/CuTeDSL/base_dsl/ast_preprocessor.py
View File

File diff suppressed because it is too large Load Diff

48
python/CuTeDSL/base_dsl/dsl.py
View File

@@ -164,16 +164,17 @@ def _mlir_type_to_numpy_type(type):
def is_dynamic_expression(value):
"""
Check if the value is an MLIR's SSA value.
Given the `value`, check if itself is an IR value or recursively go through it to check if it contains IR value
"""
# Case 1: If the value has MLIR's SSA value, return True
# Case 2: If the value supports __extract_mlir_values__ then it's possible to get SSA value
return (
isinstance(value, ir.Value)
or hasattr(value, "__extract_mlir_values__")
or len(extract_mlir_values(value)) > 0
)
if isinstance(value, (tuple, list)):
for x in value:
if is_dynamic_expression(x):
return True
elif isinstance(value, (ir.Value, ir.BlockArgumentList)) or hasattr(
value, "__extract_mlir_values__"
):
return True
return False
def extract_mlir_values(obj):
"""
@@ -726,6 +727,7 @@ class BaseDSL:
)
jit_arg_types, jit_arg_attrs, jit_exec_args = [], [], []
jit_adapted_args = []
default_attr = ir.DictAttr.get({})
input_args = [*args, *kwargs.values()]
@@ -759,7 +761,9 @@ class BaseDSL:
# If not any known type, try JIT argument adapter
# to convert the argument
adapter = JitArgAdapterRegistry.get_registered_adapter(type(arg))
arg = adapter(arg) if adapter else arg
if adapter:
arg = adapter(arg)
jit_adapted_args.append(arg)
if is_host:
jit_exec_arg.extend(get_c_pointers(arg))
@@ -798,14 +802,14 @@ class BaseDSL:
jit_arg_types.extend(jit_arg_type)
jit_arg_attrs.extend(jit_arg_attr)
return jit_exec_args, jit_arg_types, jit_arg_attrs
return jit_exec_args, jit_arg_types, jit_arg_attrs, jit_adapted_args
def generate_mlir_function_types(
self, func, function_name, input_args, kwargs, args_spec: inspect.FullArgSpec
):
"""Convert input arguments to MLIR function signature also convert numpy arrays to memref."""
exe_args, types, _ = self._generate_jit_func_args(
exe_args, types, attrs, adapted_args = self._generate_jit_func_args(
func, function_name, input_args, kwargs, args_spec, is_host=True
)
@@ -816,7 +820,7 @@ class BaseDSL:
types
), "expects the same number of arguments and function parameters"
return exe_args, types
return exe_args, types, adapted_args
@dataclass
class LaunchConfig:
@@ -1158,7 +1162,7 @@ class BaseDSL:
"""Generate MLIR module and compile iself.T_provider."""
with ir.Context(), ir.Location.unknown():
# Convert input arguments to MLIR arguments
exe_args, func_types = self.generate_mlir_function_types(
exe_args, func_types, adapted_args = self.generate_mlir_function_types(
funcBody, function_name, args, kwargs, args_spec
)
@@ -1476,7 +1480,7 @@ class BaseDSL:
if self.device_compilation_only:
return kernel_operands, kernel_arg_types, kernel_arg_attrs
kernel_operands, kernel_arg_types, kernel_arg_attrs = (
kernel_operands, kernel_arg_types, kernel_arg_attrs, _ = (
self._generate_jit_func_args(
kernel_func, kernel_name, args, kwargs, args_spec, is_host=False
)
@@ -1586,12 +1590,14 @@ class BaseDSL:
if self.device_compilation_only:
log().debug("Generating cuda-python arguments")
# Convert input arguments to MLIR arguments
self.exe_args, kernel_types = self.generate_mlir_function_types(
funcBody,
kernel_name,
canonicalized_args,
canonicalized_kwargs,
args_spec,
self.exe_args, kernel_types, _ = (
self.generate_mlir_function_types(
funcBody,
kernel_name,
canonicalized_args,
canonicalized_kwargs,
args_spec,
)
)
helper = kernelGenHelper()

6
python/CuTeDSL/base_dsl/env_manager.py
View File

@@ -78,11 +78,9 @@ def detect_gpu_arch(prefix):
major, minor = arch
suffix = ""
if major >= 9 and minor >= 0:
if major >= 9:
suffix = "a"
elif minor != 0:
# e.g sm_86, belong with sm_80 family
minor = 0
return f"sm_{major}{minor}{suffix}"

51
python/CuTeDSL/base_dsl/jit_executor.py
View File

@@ -12,23 +12,24 @@
"""
This module provides jit executor related classes
"""
import io
import inspect
import ctypes
import numpy as np
import inspect
import io
from typing import get_origin
import numpy as np
# MLIR modules imports
from .._mlir import ir
# Local modules imports
from .utils.timer import timer
from .utils.logger import log
from . import typing as t
from .common import DSLRuntimeError
from .runtime import cuda as cuda_helpers
from .runtime.jit_arg_adapters import JitArgAdapterRegistry, is_arg_spec_constexpr
from .typing import get_c_pointers
from . import typing as t
# MLIR modules imports
from .._mlir import ir
from .utils.logger import log
from .utils.timer import timer
class CudaSingleModule:
@@ -64,6 +65,7 @@ class JitExecutor:
self.args_spec = args_spec
self.function_name = function_name
if args_spec is not None:
self.original_args_spec = args_spec
self.args_spec = self.filter_runtime_arg_spec(args_spec)
# cuda kernels
self.cuda_modules = cuda_modules
@@ -135,6 +137,29 @@ class JitExecutor:
for module in set(cuda_modules):
cuda_helpers.unload_cubin_module(module)
def get_constexpr_args(self) -> list[dict[str, int | str]]:
"""
This function returns the constexpr args that have been pruned from the original function signature.
The return type is a list of dicts, each dict contains the argument index (argument_index) and argument name (argument_name).
:return: list of dicts, each dict contains the argument index (argument_index) and argument name (argument_name).
:rtype: list[dict[str, int | str]]
"""
if self.original_args_spec is None:
return list()
constexpr_args = list()
for i, arg_name in enumerate(self.original_args_spec.args):
if arg_name not in self.args_spec.args:
constexpr_args.append({"argument_index": i, "argument_name": arg_name})
if self.original_args_spec.kwonlyargs:
for kwarg in self.original_args_spec.kwonlyargs:
if kwarg not in self.args_spec.kwonlyargs:
constexpr_args.append(
{"argument_index": None, "argument_name": kwarg}
)
return constexpr_args
def generate_execution_args(self, args, kwargs, args_spec: inspect.FullArgSpec):
"""
This function is the prune version of `generate_mlir_function_types` which only generates execution args
@@ -175,6 +200,7 @@ class JitExecutor:
)
exe_args = []
adapted_args = []
input_args = rectified_args + list(rectified_kwargs.values())
input_arg_names = args_spec.args + args_spec.kwonlyargs
for arg, arg_name in zip(input_args, input_arg_names):
@@ -193,13 +219,16 @@ class JitExecutor:
adapter = JitArgAdapterRegistry.get_registered_adapter(type(arg))
if adapter:
arg = adapter(arg)
adapted_args.append(arg)
exe_args.extend(get_c_pointers(arg))
return exe_args
return exe_args, adapted_args
def __call__(self, *args, **kwargs):
exe_args = self.generate_execution_args(args, kwargs, self.args_spec)
exe_args, adapted_args = self.generate_execution_args(
args, kwargs, self.args_spec
)
self.run_compiled_program(exe_args)

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