Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -0,0 +1,584 @@
|
||||
# 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 helper functions that are generated by the preprocessor.
|
||||
The preprocessor read through python's ast and changes the input code.
|
||||
"""
|
||||
|
||||
from typing import Callable, Iterator, Optional, overload
|
||||
|
||||
from .utils.logger import log
|
||||
from .common import *
|
||||
|
||||
from ._mlir_helpers.arith import ArithValue
|
||||
|
||||
class Executor:
|
||||
"""
|
||||
The Executor class handles dynamic and compile-time (constexpr) execution
|
||||
of "for" loops and "if-else-elif" statements.
|
||||
|
||||
Methods:
|
||||
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.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._is_dynamic_expression = None
|
||||
self._loop_execute_range_dynamic = None
|
||||
self._if_dynamic = None
|
||||
self._while_dynamic = None
|
||||
|
||||
def set_functions(
|
||||
self,
|
||||
is_dynamic_expression: Callable,
|
||||
loop_execute_range_dynamic: Callable,
|
||||
if_dynamic: Callable,
|
||||
while_dynamic: Callable,
|
||||
):
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def convert_to_list(x):
|
||||
"""This function is used to convert x to a list.
|
||||
If x is None, return an empty list.
|
||||
If x is not a list, return a list containing x.
|
||||
Otherwise, return x itself.
|
||||
"""
|
||||
if x is None:
|
||||
return []
|
||||
if not isinstance(x, list):
|
||||
return [x]
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def converge_ret_val(res):
|
||||
"""This function is used to converge res (the return value) of the function.
|
||||
If res is None, return None.
|
||||
If res is a list and has only one element, return the element.
|
||||
Otherwise, return res itself.
|
||||
"""
|
||||
if res is None:
|
||||
return res
|
||||
elif isinstance(res, list) and len(res) == 1:
|
||||
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().info("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,
|
||||
start: int,
|
||||
stop: int,
|
||||
step: int,
|
||||
used_args: list,
|
||||
iter_args: list,
|
||||
):
|
||||
log().info("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):
|
||||
log().debug("i [%s] iter_args [%s]", i, iter_args)
|
||||
loop_results = func(i, *used_args, *loop_results)
|
||||
log().debug("loop_results [%s]", loop_results)
|
||||
if loop_results is None:
|
||||
loop_results = []
|
||||
if not isinstance(loop_results, list):
|
||||
loop_results = [loop_results]
|
||||
|
||||
log().debug("done loop_results [%s]", loop_results)
|
||||
return Executor.converge_ret_val(loop_results)
|
||||
|
||||
def for_execute(
|
||||
self,
|
||||
func,
|
||||
start,
|
||||
stop,
|
||||
step,
|
||||
used_args=[],
|
||||
iter_args=[],
|
||||
iter_arg_names=[],
|
||||
unroll=-1,
|
||||
unroll_full=False,
|
||||
is_range_constexpr=None,
|
||||
):
|
||||
assert (
|
||||
self._loop_execute_range_dynamic and self._is_dynamic_expression
|
||||
), "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(
|
||||
func,
|
||||
start,
|
||||
stop,
|
||||
step,
|
||||
used_args,
|
||||
iter_args,
|
||||
iter_arg_names,
|
||||
unroll,
|
||||
unroll_full,
|
||||
)
|
||||
|
||||
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,
|
||||
then_block: Callable,
|
||||
else_block: Optional[Callable] = None,
|
||||
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
|
||||
)
|
||||
|
||||
# MLIR generation
|
||||
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,
|
||||
while_before_block: Callable,
|
||||
while_after_block: Callable,
|
||||
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
|
||||
)
|
||||
|
||||
# MLIR generation
|
||||
return self.while_dynamic(
|
||||
while_before_block,
|
||||
while_after_block,
|
||||
used_args,
|
||||
yield_args,
|
||||
yield_arg_names,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Decorator
|
||||
# =============================================================================
|
||||
|
||||
executor = Executor()
|
||||
|
||||
|
||||
def loop_selector(
|
||||
start,
|
||||
stop,
|
||||
step,
|
||||
used_args=[],
|
||||
iter_args=[],
|
||||
iter_arg_names=[],
|
||||
unroll=-1,
|
||||
unroll_full=False,
|
||||
constexpr=None,
|
||||
):
|
||||
log().info(
|
||||
"start [%s] stop [%s] step [%s] used_args [%s] iter_args [%s] unroll [%s] unroll_full [%s] constexpr [%s]",
|
||||
start,
|
||||
stop,
|
||||
step,
|
||||
used_args,
|
||||
iter_args,
|
||||
unroll,
|
||||
unroll_full,
|
||||
constexpr,
|
||||
)
|
||||
from .typing import Integer, Numeric
|
||||
|
||||
def _maybe_upcast(value):
|
||||
if isinstance(value, Integer):
|
||||
value = value.ir_value()
|
||||
|
||||
return value
|
||||
|
||||
start = _maybe_upcast(start)
|
||||
stop = _maybe_upcast(stop)
|
||||
step = _maybe_upcast(step)
|
||||
|
||||
def ir_loop(func):
|
||||
return executor.for_execute(
|
||||
func,
|
||||
start,
|
||||
stop,
|
||||
step,
|
||||
used_args,
|
||||
iter_args,
|
||||
iter_arg_names,
|
||||
unroll,
|
||||
unroll_full,
|
||||
constexpr,
|
||||
)
|
||||
|
||||
return ir_loop
|
||||
|
||||
|
||||
def if_selector(pred, used_args=[], yield_args=[]):
|
||||
log().info("pred [%s] used_args [%s] yield_args [%s]", pred, used_args, yield_args)
|
||||
# Handle Numeric types here?
|
||||
|
||||
from .typing import Numeric
|
||||
|
||||
if isinstance(pred, Numeric):
|
||||
pred = pred.value
|
||||
|
||||
def ir_loop(func):
|
||||
return func(pred, *used_args, *yield_args)
|
||||
|
||||
return ir_loop
|
||||
|
||||
|
||||
def while_selector(pred, used_args=[], yield_args=[]):
|
||||
def ir_while_loop(func):
|
||||
return func(pred, *used_args, *yield_args)
|
||||
|
||||
return ir_while_loop
|
||||
|
||||
|
||||
def while_executor(
|
||||
pred,
|
||||
while_before_block: Callable,
|
||||
while_after_block: Callable,
|
||||
used_args=[],
|
||||
yield_args=[],
|
||||
yield_arg_names=[],
|
||||
constexpr=None,
|
||||
):
|
||||
return executor.while_execute(
|
||||
pred,
|
||||
while_before_block,
|
||||
while_after_block,
|
||||
used_args,
|
||||
yield_args,
|
||||
yield_arg_names,
|
||||
constexpr,
|
||||
)
|
||||
|
||||
|
||||
def if_executor(
|
||||
pred,
|
||||
then_block: Callable,
|
||||
else_block: Optional[Callable] = None,
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Range
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class range_dynamic:
|
||||
@overload
|
||||
def __new__(cls, stop, unroll=0, unroll_full=False):
|
||||
pass
|
||||
|
||||
@overload
|
||||
def __new__(cls, start, stop, step, unroll=0, unroll_full=False):
|
||||
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`.",
|
||||
)
|
||||
|
||||
def __iter__(self) -> Iterator[int]:
|
||||
current = self.start
|
||||
while current < self.stop:
|
||||
yield current
|
||||
current += self.step
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# If expressions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def const_expr(expression):
|
||||
if executor._is_dynamic_expression(expression):
|
||||
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",
|
||||
},
|
||||
)
|
||||
return expression
|
||||
|
||||
|
||||
def dynamic_expr(expression):
|
||||
raise DSLRuntimeError("dynamic_expr should be always preprocessed to IR")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Assertion & casting
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def assert_executor(test, msg=None):
|
||||
from .typing import Numeric
|
||||
|
||||
fail = False
|
||||
# Implicit convert dynamic expression to bool is not allowed
|
||||
# So here explicitly do a None check
|
||||
if test is not None and executor._is_dynamic_expression(test):
|
||||
if isinstance(test, Numeric):
|
||||
try:
|
||||
test = test.to(bool)
|
||||
except:
|
||||
fail = True
|
||||
else:
|
||||
fail = True
|
||||
|
||||
if not fail:
|
||||
assert test, msg
|
||||
else:
|
||||
raise DSLRuntimeError(
|
||||
"Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.",
|
||||
suggestion = "Please replace with runtime assert."
|
||||
)
|
||||
|
||||
|
||||
def bool_cast(value):
|
||||
if executor._is_dynamic_expression(value):
|
||||
raise DSLRuntimeError(
|
||||
"Only constexpr (Python Value) is allowed here, but got non-constexpr (IR Values) expression.",
|
||||
suggestion = "Please explicitly convert to boolean with expressions like comparision."
|
||||
)
|
||||
return bool(value)
|
||||
Reference in New Issue
Block a user