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

v4.0 update. (#2371)

Browse Source
This commit is contained in:
Junkai-Wu
2025-06-06 14:39:20 +08:00
committed by GitHub
parent 2e2af190bd
commit 8bdbfca682
254 changed files with 29751 additions and 1980 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/CuTeDSL/base_dsl/_mlir_helpers/arith.py
View File

@@ -578,7 +578,7 @@ class ArithValue(ir.Value):
# Unary operators
def __invert__(self, *, loc=None, ip=None) -> "ArithValue":
return arith.xori(self, arith.const(self.type, -1))
return arith.xori(self, arith.constant(self.type, -1))
# Bitwise operations
@_dispatch_to_rhs_r_op

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

@@ -95,7 +95,7 @@ class Executor:
unroll=bool,
unroll_full=int,
):
log().info("start [%s] stop [%s] step [%s]", start, stop, step)
log().debug("start [%s] stop [%s] step [%s]", start, stop, step)
return self._loop_execute_range_dynamic(
func,
start,
@@ -117,7 +117,7 @@ class Executor:
used_args: list,
iter_args: list,
):
log().info("start [%s] stop [%s] step [%s]", start, stop, step)
log().debug("start [%s] stop [%s] step [%s]", start, stop, step)
loop_results = iter_args
log().debug("iter_args [%s]", iter_args)
for i in range(start, stop, step):
@@ -374,7 +374,7 @@ def loop_selector(
unroll_full=False,
constexpr=None,
):
log().info(
log().debug(
"start [%s] stop [%s] step [%s] used_args [%s] iter_args [%s] unroll [%s] unroll_full [%s] constexpr [%s]",
start,
stop,
@@ -415,7 +415,7 @@ def loop_selector(
def if_selector(pred, used_args=[], yield_args=[]):
log().info("pred [%s] used_args [%s] yield_args [%s]", pred, used_args, yield_args)
log().debug("pred [%s] used_args [%s] yield_args [%s]", pred, used_args, yield_args)
# Handle Numeric types here?
from .typing import Numeric

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

@@ -248,39 +248,55 @@ class DSLPreprocessor(ast.NodeTransformer):
# Step 3. Return the transformed tree
return combined_body
def check_early_exit(self, tree):
def check_early_exit(self, tree, kind):
"""
Checks if a given region or scope in the provided Python code has early exits.
"""
class EarlyExitChecker(ast.NodeVisitor):
def __init__(self):
def __init__(self, kind):
self.has_early_exit = False
self.early_exit_node = None
self.early_exit_type = None
self.kind = kind
self.loop_nest_level = 0
# Early exit is not allowed in any level of dynamic control flow
def visit_Return(self, node):
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "return"
def visit_Break(self, node):
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "break"
def visit_Continue(self, node):
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "continue"
def visit_Raise(self, node):
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "raise"
checker = EarlyExitChecker()
checker.visit(tree)
def visit_Break(self, node):
# For break/continue in inner loops, we don't consider it as early exit
if self.loop_nest_level == 0 and self.kind != "if":
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "break"
def visit_Continue(self, node):
if self.loop_nest_level == 0 and self.kind != "if":
self.has_early_exit = True
self.early_exit_node = node
self.early_exit_type = "continue"
def visit_For(self, node):
self.loop_nest_level += 1
self.generic_visit(node)
self.loop_nest_level -= 1
def visit_While(self, node):
self.loop_nest_level += 1
self.generic_visit(node)
self.loop_nest_level -= 1
checker = EarlyExitChecker(kind)
checker.generic_visit(tree)
if not checker.has_early_exit:
return
raise DSLAstPreprocessorError(
@@ -591,7 +607,7 @@ class DSLPreprocessor(ast.NodeTransformer):
if self.is_supported_range_call(node):
constexpr_val = self.get_loop_constexpr(node)
# Check for early exit and raise exception
self.check_early_exit(node)
self.check_early_exit(node, "for")
start, stop, step = self.extract_range_args(node.iter)
unroll, unroll_full = self.extract_unroll_args(node.iter)
used_args, iter_args, flat_args = self.analyze_region_variables(
@@ -659,37 +675,42 @@ class DSLPreprocessor(ast.NodeTransformer):
snippet=ast.unparse(node),
)
test = ast.BoolOp(
op=ast.And(),
values=[
ast.Compare(
left=ast.Call(
func=ast.Name(id="type", ctx=ast.Load()),
args=[node.values[0]],
keywords=[],
def short_circuit_eval(value, short_circuit_value):
return ast.BoolOp(
op=ast.And(),
values=[
ast.Compare(
left=ast.Call(
func=ast.Name(id="type", ctx=ast.Load()),
args=[value],
keywords=[],
),
ops=[ast.Eq()],
comparators=[ast.Name(id="bool", ctx=ast.Load())],
),
ops=[ast.Eq()],
comparators=[ast.Name(id="bool", ctx=ast.Load())],
),
ast.Compare(
left=node.values[0],
ops=[ast.Eq()],
comparators=[short_circuit_value],
),
],
)
return ast.copy_location(
ast.IfExp(
ast.Compare(
left=value,
ops=[ast.Eq()],
comparators=[short_circuit_value],
),
],
)
lhs = node.values[0]
for i in range(1, len(node.values)):
test = short_circuit_eval(lhs, short_circuit_value)
lhs = ast.IfExp(
test=test,
body=node.values[0],
body=lhs,
orelse=ast.Call(
func=helper_func,
args=node.values,
args=[lhs, node.values[i]],
keywords=[],
),
),
node,
)
)
return ast.copy_location(lhs, node)
def visit_UnaryOp(self, node):
# Visit child nodes first
@@ -916,7 +937,7 @@ class DSLPreprocessor(ast.NodeTransformer):
return node
# Check for early exit and raise exception
self.check_early_exit(node)
self.check_early_exit(node, "while")
used_args, yield_args, flat_args = self.analyze_region_variables(
node, active_symbols
@@ -1021,7 +1042,7 @@ class DSLPreprocessor(ast.NodeTransformer):
return node
# Check for early exit and raise exception
self.check_early_exit(node)
self.check_early_exit(node, "if")
used_args, yield_args, flat_args = self.analyze_region_variables(
node, active_symbols

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

@@ -566,7 +566,9 @@ class BaseDSL:
log().debug("Processing [%d] Argument [%s : %s]", i, arg_name, arg_spec)
# Implicit cast to NumericMeta
if isinstance(arg_spec, t.NumericMeta):
if isinstance(arg_spec, t.NumericMeta) and not isinstance(
arg, arg_spec
):
arg = t.cast(arg, arg_spec)
ir_arg, iv_block_args = (
@@ -589,15 +591,17 @@ class BaseDSL:
self.log_additions(ir_arg)
ir_args.extend(ir_arg)
return ir_args
return ir_args, iv_block_args
fop_args = list(fop.regions[0].blocks[0].arguments)
ir_args = gen_exec_args(args, args_spec.args, args_spec.annotations, fop_args)
ir_kwargs = gen_exec_args(
ir_args, iv_block_args = gen_exec_args(
args, args_spec.args, args_spec.annotations, fop_args
)
ir_kwargs, _ = gen_exec_args(
[kwargs[arg] for arg in args_spec.kwonlyargs],
args_spec.kwonlyargs,
args_spec.annotations,
fop_args[len(ir_args) :],
fop_args[iv_block_args:],
)
ir_kwargs = {k: v for k, v in zip(args_spec.kwonlyargs, ir_kwargs)}
@@ -716,8 +720,10 @@ class BaseDSL:
assert len(args) == len(args_spec.args) and len(kwargs) == len(
args_spec.kwonlyargs
), f"Input args {len(args)=} and kwargs {len(kwargs)=} must match arg_spec.args "
f"{len(args_spec.args)=} and arg_spec.kwonlyargs {len(args_spec.kwonlyargs)=}"
), (
f"Input args {len(args)=} and kwargs {len(kwargs)=} must match arg_spec.args "
f"{len(args_spec.args)=} and arg_spec.kwonlyargs {len(args_spec.kwonlyargs)=}"
)
jit_arg_types, jit_arg_attrs, jit_exec_args = [], [], []
default_attr = ir.DictAttr.get({})
@@ -729,7 +735,7 @@ class BaseDSL:
log().debug("Processing [%d] Argument [%s : %s]", i, arg_name, spec_ty)
# Implicitly convert into Numeric type if possible
if isinstance(spec_ty, t.NumericMeta):
if isinstance(spec_ty, t.NumericMeta) and not isinstance(arg, spec_ty):
arg = t.cast(arg, spec_ty)
# Type safety check

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

@@ -141,33 +141,60 @@ class JitExecutor:
to get rid of mlir context.
"""
# Process positional arguments with defaults
rectified_args = list(args)
if args_spec.defaults and len(args) < len(args_spec.args):
rectified_args.extend(args_spec.defaults[len(args) - len(args_spec.args) :])
for k, v in kwargs.items():
if k in args_spec.args:
idx = args_spec.args.index(k)
if idx < len(rectified_args):
rectified_args[idx] = v
else:
rectified_args.append(v)
# Process keyword arguments
rectified_kwargs = {k: v for k, v in kwargs.items() if k not in args_spec.args}
if args_spec.kwonlydefaults and len(rectified_kwargs) < len(
args_spec.kwonlyargs
):
rectified_kwargs.update(args_spec.kwonlydefaults)
# args/kwargs must match arg_specs
# No canonicalization of args/kwargs to avoid extra latency
if len(args) != len(args_spec.args) or len(kwargs) != len(args_spec.kwonlyargs):
if len(rectified_args) != len(args_spec.args) or len(rectified_kwargs) != len(
args_spec.kwonlyargs
):
raise DSLRuntimeError(
"input args/kwargs length does not match runtime function signature!",
context={
"input args length": len(args),
"input kwargs length": len(kwargs),
"input args length": len(rectified_args),
"input kwargs length": len(rectified_kwargs),
"function signature args length": len(args_spec.args),
"function signature kwonlyargs length": len(args_spec.kwonlyargs),
},
)
exe_args = []
input_args = [*args, *kwargs.values()]
input_arg_names = [*args_spec.args, *args_spec.kwonlyargs]
for i, arg in enumerate(input_args):
arg_type = args_spec.annotations.get(input_arg_names[i], None)
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):
# short-cut for args already converted
if hasattr(arg, "__c_pointers__"):
exe_args.extend(arg.__c_pointers__())
continue
arg_type = args_spec.annotations.get(arg_name, None)
# Implicit cast to NumericMeta
if isinstance(arg_type, t.NumericMeta):
arg = t.cast(arg, arg_type)
else:
# If not any known type, try registered adapter to do the conversion
adapter = JitArgAdapterRegistry.get_registered_adapter(type(arg))
if adapter:
arg = adapter(arg)
# If not any known type, try registered adapter to do the conversion
adapter = JitArgAdapterRegistry.get_registered_adapter(type(arg))
adapted_arg = adapter(arg) if adapter else arg
exe_args.extend(get_c_pointers(adapted_arg))
exe_args.extend(get_c_pointers(arg))
return exe_args

2
python/CuTeDSL/base_dsl/runtime/cuda.py
View File

@@ -457,7 +457,7 @@ class StreamAdapter:
def __init__(self, arg):
self._arg = arg
self._c_pointer = ctypes.cast(self._arg.getPtr(), ctypes.c_void_p)
self._c_pointer = self._arg.getPtr()
def __new_from_mlir_values__(self, values):
assert len(values) == 1

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

@@ -629,7 +629,7 @@ def _binary_op_type_promote(a, b, promote_bool: bool = False):
b_type = b.dtype
# Early return for same types (except when they're bools that need promotion)
if a_type == b_type and not (promote_bool and a_type.width == 1):
if a_type == b_type and not (promote_bool and a_type is Boolean):
return a, b, a_type
# Handle floating point promotions
@@ -1315,10 +1315,7 @@ class Integer(Numeric, metaclass=IntegerMeta, mlir_type=T.i32, is_abstract=True)
def __invert__(self, *, loc=None, ip=None):
res_type = type(self)
# Create a constant of -1 (all bits set to 1) of the same type as value
all_ones = arith.constant(res_type.mlir_type, -1)
# XOR with -1 gives us bitwise NOT
return res_type(arith.xori(self.ir_value(), all_ones, loc=loc, ip=ip))
return res_type(self.ir_value(loc=loc, ip=ip).__invert__(loc=loc, ip=ip))
def __lshift__(self, other, *, loc=None, ip=None):
return _binary_op(operator.lshift)(self, other, loc=loc, ip=ip)
@@ -1457,18 +1454,14 @@ class Boolean(Integer, metaclass=IntegerMeta, width=1, signed=True, mlir_type=T.
- Converted using Python's bool() function
- Example: Boolean(1) -> True, Boolean(0) -> False
2. Boolean:
- Direct value assignment
- Example: Boolean(Boolean(True)) -> True
2. Numeric:
- Uses the Numeric.value to construct Boolean recursively
3. Numeric:
- Uses the __dsl_bool__ method of the Numeric type
4. MLIR Value with IntegerType:
3. MLIR Value with IntegerType:
- If width is 1: Direct assignment
- Otherwise: Compares with 0 using arith.cmpi
5. MLIR Value with FloatType:
4. MLIR Value with FloatType:
- Compares with 0.0 using arith.cmpf
- Uses unordered comparison to handle NaN values
"""
@@ -1479,19 +1472,35 @@ class Boolean(Integer, metaclass=IntegerMeta, width=1, signed=True, mlir_type=T.
value = None
if isinstance(a, (bool, int, float)):
value = bool(a)
elif isinstance(a, Boolean):
value = a.value
elif isinstance(a, Numeric):
value = a.__dsl_bool__(loc=loc, ip=ip)
Boolean.__init__(self, a.value, loc=loc, ip=ip)
return
elif isinstance(a, ArithValue):
if a.type == T.bool():
value = a
else:
value = a != arith_helper.const(0, a.type)
value = a != arith_helper.const(0, a.type, loc=loc, ip=ip)
if value is None:
raise DSLRuntimeError(f"Cannot convert {a} to Boolean")
super().__init__(value, loc=loc, ip=ip)
self._value_int8 = None
def ir_value_int8(self, *, loc=None, ip=None):
"""
Returns int8 ir value of Boolean.
When we need to store Boolean tensor element, use ir_value_int8().
:param loc: Source location information, defaults to None
:type loc: Optional[Location], optional
:param ip: Insertion point for MLIR operations, defaults to None
:type ip: Optional[InsertionPoint], optional
:return: The int8 value of this Boolean
:rtype: ir.Value
"""
if self._value_int8 is not None:
return self._value_int8
self._value_int8 = Int8(self.value, loc=loc, ip=ip).ir_value()
return self._value_int8
def __neg__(self, *, loc=None, ip=None):
"""Negation operator is not supported for boolean type.