This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d0002f3c6a [RELAX] Unify call_tir primitive arguments (#20009)
d0002f3c6a is described below
commit d0002f3c6a63b39477d2c129e805e8267f79c43f
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Jul 16 05:02:36 2026 +0800
[RELAX] Unify call_tir primitive arguments (#20009)
---
python/tvm/relax/block_builder.py | 12 +-
python/tvm/relax/dpl/pattern.py | 10 +-
python/tvm/relax/frontend/nn/op.py | 23 ++--
python/tvm/relax/op/base.py | 37 ++----
python/tvm/relax/op/distributed/distributed.py | 14 +--
python/tvm/relax/script/builder/distributed/ir.py | 14 +--
python/tvm/relax/script/builder/ir.py | 4 +-
.../tvm/relax/transform/legalize_ops/manipulate.py | 8 +-
python/tvm/relax/utils.py | 18 +--
.../transform/lower_global_view_to_local_view.cc | 24 +++-
src/relax/ir/dataflow_pattern.cc | 6 +-
src/relax/op/distributed/distributed.cc | 17 +--
src/relax/op/op.cc | 133 ++++-----------------
src/relax/script/builder/distributed.cc | 12 +-
src/relax/script/printer/call.cc | 7 +-
src/relax/transform/call_tir_rewrite.cc | 11 +-
src/relax/transform/fold_constant.cc | 3 -
src/relax/transform/fuse_tir.cc | 103 +++++-----------
...ributed_transform_lower_global_to_local_view.py | 4 -
...est_distributed_transform_propagate_sharding.py | 16 +--
tests/python/relax/test_analysis_well_formed.py | 70 +++++++++++
tests/python/relax/test_blockbuilder_emit_te.py | 12 +-
tests/python/relax/test_dataflow_pattern.py | 8 +-
tests/python/relax/test_frontend_nn_op.py | 17 +--
tests/python/relax/test_op_index.py | 9 +-
tests/python/relax/test_transform.py | 36 ++++++
tests/python/relax/test_transform_fold_constant.py | 10 +-
tests/python/relax/test_transform_fuse_tir.py | 75 ++++++++----
.../relax/test_transform_lazy_transform_params.py | 10 +-
.../test_transform_legalize_ops_create_datatype.py | 12 +-
.../test_transform_legalize_ops_distributed.py | 4 +-
..._transform_legalize_ops_index_linear_algebra.py | 4 +-
.../test_transform_legalize_ops_manipulate.py | 4 +-
.../python/relax/test_transform_legalize_ops_nn.py | 4 +-
.../relax/test_transform_lift_transform_params.py | 16 +--
.../test_transform_rewrite_dataflow_reshape.py | 10 +-
tests/python/relax/test_tvmscript_parser.py | 4 +-
tests/python/relax/test_tvmscript_printer_relax.py | 8 +-
38 files changed, 344 insertions(+), 445 deletions(-)
diff --git a/python/tvm/relax/block_builder.py
b/python/tvm/relax/block_builder.py
index 3a2cd4178d..d67c70e0cb 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -361,13 +361,13 @@ class BlockBuilder(Object):
"""
primfunc_name = kwargs.pop("primfunc_name_hint", None)
- tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func,
*args, **kwargs)
+ tir_func, call_args, output_ty = gen_call_tir_inputs(func, *args,
**kwargs)
if not primfunc_name:
primfunc_name = func.__name__
gvar = self.add_func(tir_func, primfunc_name)
- return call_tir(gvar, call_args, output_ty, tir_vars)
+ return call_tir(gvar, call_args, output_ty)
def call_te_with_grad(
self,
@@ -413,7 +413,7 @@ class BlockBuilder(Object):
"""
primfunc_name = kwargs.pop("primfunc_name_hint", None)
- tir_func, call_args, output_ty, tir_vars = gen_call_tir_inputs(func,
*args, **kwargs)
+ tir_func, call_args, output_ty = gen_call_tir_inputs(func, *args,
**kwargs)
if te_grad_kwargs is None:
te_grad_kwargs = {}
@@ -422,9 +422,7 @@ class BlockBuilder(Object):
primfunc_name = func.__name__
gvar = self.add_func(tir_func, primfunc_name)
- return call_tir_with_grad(
- gvar, call_args, output_ty, te_grad_name, te_grad_kwargs, tir_vars
- )
+ return call_tir_with_grad(gvar, call_args, output_ty, te_grad_name,
te_grad_kwargs)
def emit_te(self, func: Callable, *args: Any, **kwargs: Any) -> Var:
"""Emit a call node according to the te function.
@@ -541,7 +539,7 @@ class BlockBuilder(Object):
def rx_func(x: Tensor((n,), "float32"), y: Tensor(((n + 1),),
"float32"))
-> Tensor(None, "float32", ndim=-1):
# block 0
- gv = relax.call_tir(te_func, (y,), R.Tensor((n + 1,),
"float32"), (n,))
+ gv = relax.call_tir(te_func, (y, n), R.Tensor((n + 1,),
"float32"))
return gv
"""
name_hint = kwargs.pop("name_hint", "")
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
index 648de6ebf1..caa04bb8f0 100644
--- a/python/tvm/relax/dpl/pattern.py
+++ b/python/tvm/relax/dpl/pattern.py
@@ -864,23 +864,19 @@ def is_shape(shape: list[tvm.ir.Expr]) ->
"PrimArrPattern":
def _is_call_tir(
func_pattern: DFPattern,
args: list | tuple | TuplePattern = None,
- tir_vars: DFPattern | None = None,
) -> CallPattern:
if args is None:
args = wildcard()
elif isinstance(args, list | tuple):
args = TuplePattern(args)
- if tir_vars is None:
- return is_op("relax.call_tir")(func_pattern, args,
add_constraint=False)
- return is_op("relax.call_tir")(func_pattern, args, tir_vars,
add_constraint=False)
+ return is_op("relax.call_tir")(func_pattern, args, add_constraint=False)
# Todo(relax-team): Dataflow pattern for Type, and match out_ty
def is_call_tir(
func_name: str,
args: list | tuple | TuplePattern = None,
- tir_vars: DFPattern | None = None,
) -> CallPattern:
"""
Syntax sugar for creating a CallPattern for call_tir that calls an
function through global var.
@@ -891,15 +887,13 @@ def is_call_tir(
Name of the CPS function to call.
args : Union[List[DFPattern], Tuple[DFPattern]], optional
Arguments in expected call_packed, by default None meaning arbitrary
(number of) arguments
- tir_vars : Optional[DFPattern]
- Pattern to match the tuple of integers that are unpacked when calling
the tirx func.
Returns
-------
CallPattern
The resulting CallPattern
"""
func_pattern = GlobalVarPattern(func_name)
- return _is_call_tir(func_pattern, args, tir_vars)
+ return _is_call_tir(func_pattern, args)
def _is_call_dps_packed(
diff --git a/python/tvm/relax/frontend/nn/op.py
b/python/tvm/relax/frontend/nn/op.py
index a12c96cdbe..13a93097ff 100644
--- a/python/tvm/relax/frontend/nn/op.py
+++ b/python/tvm/relax/frontend/nn/op.py
@@ -2069,15 +2069,17 @@ def tensor_ir_op(
"""
from tvm import relax as rx # pylint: disable=import-outside-toplevel
- call_tir_args, tir_vars = [], []
+ call_tir_args = []
if not isinstance(args, tuple | list):
args = [args]
for arg in args:
if isinstance(arg, Tensor):
call_tir_args.append(arg._expr)
- elif isinstance(arg, rx.ShapeExpr) or tvm.ir.is_prim_expr(arg):
- tir_vars.append(arg)
+ elif isinstance(arg, rx.ShapeExpr):
+ call_tir_args.extend(arg.values)
+ elif tvm.ir.is_prim_expr(arg):
+ call_tir_args.append(arg)
else:
raise TypeError(
"Unsupported type: tensor_ir_op args expect Tensor or
ShapeExpr or Expr,"
@@ -2092,11 +2094,8 @@ def tensor_ir_op(
bb = BlockBuilder.current()
global_var = bb.add_func(func, name_hint)
- if len(tir_vars) == 0:
- tir_vars = None
-
return wrap_nested(
- bb.emit(rx.call_tir(global_var, call_tir_args, out_ty,
tir_vars=tir_vars)),
+ bb.emit(rx.call_tir(global_var, call_tir_args, out_ty)),
name=name_hint,
)
@@ -2139,15 +2138,17 @@ def tensor_ir_inplace_op(
"""
from tvm import relax as rx # pylint: disable=import-outside-toplevel
- call_tir_args, tir_vars = [], []
+ call_tir_args = []
if not isinstance(args, tuple | list):
args = [args]
for arg in args:
if isinstance(arg, Tensor):
call_tir_args.append(arg._expr)
- elif isinstance(arg, rx.ShapeExpr) or tvm.ir.is_prim_expr(arg):
- tir_vars.append(arg)
+ elif isinstance(arg, rx.ShapeExpr):
+ call_tir_args.extend(arg.values)
+ elif tvm.ir.is_prim_expr(arg):
+ call_tir_args.append(arg)
else:
raise TypeError(
"Unsupported type: tensor_ir_inplace_op args expect Tensor or
ShapeExpr or"
@@ -2163,7 +2164,7 @@ def tensor_ir_inplace_op(
global_var = bb.add_func(func, name_hint)
return wrap_nested(
- bb.emit(rx.call_tir_inplace(global_var, call_tir_args,
inplace_indices, out_ty, tir_vars)),
+ bb.emit(rx.call_tir_inplace(global_var, call_tir_args,
inplace_indices, out_ty)),
name=name_hint,
)
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index e1a3d46002..e0e89d49bb 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -26,7 +26,7 @@ import tvm.runtime
from tvm.ir import Call
from tvm.runtime import Object, ObjectConvertible
-from ..expr import Expr, ExternFunc, GlobalVar, ShapeExpr, StringImm, Var
+from ..expr import Expr, ExternFunc, GlobalVar, StringImm, Var
from ..type import TensorType, Type
from ..utils import convert_to_expr
from . import _ffi_api
@@ -93,7 +93,6 @@ def call_tir(
gvar: GlobalVar,
args: Expr,
out_ty: TensorType | list[TensorType],
- tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output.
@@ -104,16 +103,14 @@ def call_tir(
The GlobalVar referring to a tirx PrimFunc.
args : Expr
- The input arguments.
+ The ordered tensor and primitive input arguments. These correspond
+ positionally to the leading parameters of the PrimFunc.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_tir output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
- tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
- ShapeExpr representing a tuple of integers to unpack when calling
func. Is null if not used
-
Returns
-------
ret: Call
@@ -124,10 +121,7 @@ def call_tir(
if not isinstance(out_ty, list):
out_ty = [out_ty]
- if isinstance(tir_vars, list | tuple):
- tir_vars = ShapeExpr(tir_vars)
-
- return _ffi_api.call_tir(gvar, args, out_ty, tir_vars) # type: ignore
+ return _ffi_api.call_tir(gvar, args, out_ty) # type: ignore
def call_tir_with_grad(
@@ -136,7 +130,6 @@ def call_tir_with_grad(
out_ty: TensorType | list[TensorType],
te_grad_name: str,
te_grad_kwargs: dict[str, Object] | None = None,
- tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output. This intrinsic will bind a te
gradient function
@@ -149,7 +142,8 @@ def call_tir_with_grad(
The GlobalVar referring to a tirx PrimFunc.
args : Expr
- The input arguments.
+ The ordered tensor and primitive input arguments. These correspond
+ positionally to the leading parameters of the PrimFunc.
out_ty : Union[TensorType, List[TensorType]]
The type information of the call_tir_with_grad output.
@@ -164,9 +158,6 @@ def call_tir_with_grad(
The keyword arguments passed to the te gradient function.
Optionally provided as a keyword argument. Default: {}.
- tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
- ShapeExpr representing a tuple of integers to unpack when calling
func. Is null if not used
-
Returns
-------
ret: Call
@@ -177,14 +168,11 @@ def call_tir_with_grad(
if not isinstance(out_ty, list):
out_ty = [out_ty]
- if isinstance(tir_vars, list | tuple):
- tir_vars = ShapeExpr(tir_vars)
-
if te_grad_kwargs is None:
te_grad_kwargs = {}
return _ffi_api.call_tir_with_grad( # type: ignore
- gvar, args, out_ty, te_grad_name, te_grad_kwargs, tir_vars
+ gvar, args, out_ty, te_grad_name, te_grad_kwargs
)
@@ -193,7 +181,6 @@ def call_tir_inplace(
args: Expr,
inplace_indices: int | list[int],
out_ty: TensorType | list[TensorType],
- tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a TIR PrimFunc and return the result, doing the specified
computations in-place
@@ -214,7 +201,8 @@ def call_tir_inplace(
The GlobalVar referring to a TIR PrimFunc.
args : Expr
- The input arguments.
+ The ordered tensor and primitive input arguments. These correspond
+ positionally to the leading parameters of the PrimFunc.
inplace_indices : Union[int, List[int]]
Specify which arguments should be used for in-place computations.
@@ -230,9 +218,6 @@ def call_tir_inplace(
Each one denotes the type information of a returned tensor.
If a list of `TensorType` is given, the result will be a tuple of
`TensorType`.
- tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
- ShapeExpr representing a tuple of integers to unpack when calling
func. Is null if not used
-
Returns
-------
ret: Call
@@ -246,15 +231,11 @@ def call_tir_inplace(
if not isinstance(out_ty, list):
out_ty = [out_ty]
- if isinstance(tir_vars, list | tuple):
- tir_vars = ShapeExpr(tir_vars)
-
return _ffi_api.call_tir_inplace( # type: ignore
gvar,
args,
inplace_indices,
out_ty,
- tir_vars,
)
diff --git a/python/tvm/relax/op/distributed/distributed.py
b/python/tvm/relax/op/distributed/distributed.py
index e39b227669..38d79ffc84 100644
--- a/python/tvm/relax/op/distributed/distributed.py
+++ b/python/tvm/relax/op/distributed/distributed.py
@@ -20,7 +20,7 @@
from tvm.ir import Call
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
-from ...expr import Expr, GlobalVar, ShapeExpr
+from ...expr import Expr, GlobalVar
from ...expr import Tuple as RxTuple
from ...utils import convert_to_expr
from . import _ffi_api
@@ -69,7 +69,6 @@ def call_tir_local_view(
gvar: GlobalVar,
args: Expr,
out_ty: DTensorType | list[DTensorType],
- tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""
Call a tirx.prim_func and return the output. The prim_func should be a
worker-local function
@@ -82,16 +81,14 @@ def call_tir_local_view(
The GlobalVar referring to a tirx PrimFunc.
args : Expr
- The input arguments.
+ The ordered distributed-tensor and primitive input arguments. These
+ correspond positionally to the leading parameters of the PrimFunc.
out_ty : Union[DTensorType, List[DTensorType]]
The type information of the call_tir output.
It should be a single or a list of DTensorType. Each one denotes the
type information of a returned tensor.
- tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
- ShapeExpr representing a tuple of integers to unpack when calling
func. Is null if not used
-
Returns
-------
ret: Call
@@ -105,10 +102,7 @@ def call_tir_local_view(
if not isinstance(out_ty, list):
out_ty = [out_ty]
- if isinstance(tir_vars, list | tuple):
- tir_vars = ShapeExpr(tir_vars)
-
- return _ffi_api.call_tir_local_view(gvar, args, out_ty, tir_vars) # type:
ignore
+ return _ffi_api.call_tir_local_view(gvar, args, out_ty) # type: ignore
def redistribute_replica_to_shard(input: Expr, num_workers: int, axis: int) ->
Expr:
diff --git a/python/tvm/relax/script/builder/distributed/ir.py
b/python/tvm/relax/script/builder/distributed/ir.py
index 82d5e9805e..738199ab25 100644
--- a/python/tvm/relax/script/builder/distributed/ir.py
+++ b/python/tvm/relax/script/builder/distributed/ir.py
@@ -28,7 +28,7 @@ import tvm
from tvm import base as _base
from tvm.ir import Call
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
-from tvm.relax.expr import Constant, Expr, ExternFunc, ShapeExpr
+from tvm.relax.expr import Constant, Expr, ExternFunc
from tvm.relax.expr import Tuple as RxTuple
from tvm.relax.op.distributed import (
annotate_sharding as _annotate_sharding,
@@ -53,7 +53,6 @@ def call_tir(
func: str | Expr,
args: Expr,
out_ty: DTensorType | list[DTensorType],
- tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""Distributed version of call_tir
@@ -63,16 +62,14 @@ def call_tir(
The destination-passing-style function, can be ExternFunc or PrimFunc.
args : Expr
- The input arguments.
+ The ordered distributed-tensor and primitive input arguments. These
+ correspond positionally to the leading parameters of the PrimFunc.
out_ty : Union[DTensorType, List[DTensorType]]
The type information of the call_tir output.
It should be a single or a list of DTensorType. Each one denotes the
type information of a returned distributed tensor.
- tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
- ShapeExpr representing a tuple of integers to unpack when calling
func. Is null if not used
-
Returns
-------
ret: Call
@@ -89,10 +86,7 @@ def call_tir(
if not isinstance(out_ty, list):
out_ty = [out_ty]
- if isinstance(tir_vars, list | tuple):
- tir_vars = ShapeExpr(tir_vars)
-
- return _ffi_api.call_tir_dist(func, args, out_ty, tir_vars) # type: ignore
+ return _ffi_api.call_tir_dist(func, args, out_ty) # type: ignore
def const(
diff --git a/python/tvm/relax/script/builder/ir.py
b/python/tvm/relax/script/builder/ir.py
index 0396ebec60..1f19fbb05b 100644
--- a/python/tvm/relax/script/builder/ir.py
+++ b/python/tvm/relax/script/builder/ir.py
@@ -582,11 +582,11 @@ def emit_te(func: Callable, *args: Any, **kwargs: Any) ->
Call:
A newly created call that calls into a tirx function.
"""
primfunc_name_hint = kwargs.pop("primfunc_name_hint", None)
- tir_func, call_args, out_ty, tir_vars = gen_call_tir_inputs(func, *args,
**kwargs)
+ tir_func, call_args, out_ty = gen_call_tir_inputs(func, *args, **kwargs)
if not primfunc_name_hint:
primfunc_name_hint = func.__name__
gvar = decl_function(primfunc_name_hint, tir_func) # type: ignore
- return call_tir(gvar, call_args, out_ty, tir_vars)
+ return call_tir(gvar, call_args, out_ty)
def emit_match_cast(value: Expr, ty: Type) -> Var:
diff --git a/python/tvm/relax/transform/legalize_ops/manipulate.py
b/python/tvm/relax/transform/legalize_ops/manipulate.py
index f94dc67501..ac72f032f9 100644
--- a/python/tvm/relax/transform/legalize_ops/manipulate.py
+++ b/python/tvm/relax/transform/legalize_ops/manipulate.py
@@ -321,7 +321,7 @@ def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
def te_layout_transform(data, name):
"""
Returns a passthrough TE compute with appropriate name. This is needed
to generate
- TIR function, output shape info, TIR vars from gen_call_tir_inputs
function.
+ TIR function and output shape info from gen_call_tir_inputs function.
"""
return te.compute(
data.shape,
@@ -353,9 +353,7 @@ def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
primfunc_name += "_with_pad"
if len(axis_separators) != 0:
primfunc_name += "_axis_separator"
- tir_func, call_args, _, tir_vars = gen_call_tir_inputs(
- te_layout_transform, call.args[0], primfunc_name
- )
+ tir_func, call_args, _ = gen_call_tir_inputs(te_layout_transform,
call.args[0], primfunc_name)
# Create TIR schedule to apply layout changes with axis separators
sch = tvm.s_tir.Schedule(tir_func)
sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
@@ -366,4 +364,4 @@ def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
output_shape = index_map.map_shape(list(call_args[0].ty.shape))
output_dtype = call_args[0].ty.dtype
output_ty = [TensorType(output_shape, output_dtype)]
- return call_tir(gvar, call_args, output_ty, tir_vars)
+ return call_tir(gvar, call_args, output_ty)
diff --git a/python/tvm/relax/utils.py b/python/tvm/relax/utils.py
index a7cc1403a8..78242bd070 100644
--- a/python/tvm/relax/utils.py
+++ b/python/tvm/relax/utils.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# ruff: noqa: F401, RUF005
+# ruff: noqa: F401
# pylint: disable=invalid-name,too-many-locals
@@ -145,7 +145,7 @@ def copy_with_new_vars(func: Function) -> Function:
def gen_call_tir_inputs(
func: Callable, *args: Any, **kwargs: Any
-) -> tuple[tirx.PrimFunc, Expr, list[TensorType], ShapeExpr | None]:
+) -> tuple[tirx.PrimFunc, Expr, list[TensorType]]:
"""Generate the inputs for call_tir according to the te function.
This function converts arguments from relax expression to te tensor,
The callback func should return a te tensor or a list of te tensors.
@@ -165,9 +165,9 @@ def gen_call_tir_inputs(
Returns
-------
- ret : Tuple[tirx.PrimFunc, Expr, List[TensorType], Optional[ShapeExpr]]
+ ret : Tuple[tirx.PrimFunc, Expr, List[TensorType]]
ret contains the inputs for call_tir, including a tirx prim_func, args,
- out_ty, and tir_vars.
+ and out_ty.
"""
tir_var_map: dict[tvm.ir.Var, tirx.Var] = {}
@@ -350,7 +350,7 @@ def gen_call_tir_inputs(
outs = [te_out] if isinstance(te_out, te_Tensor) else list(te_out)
unbound_tir_vars = _get_unbound_tir_vars([*create_primfunc_args, *outs],
extra_tir_args_list)
- inputs = [*create_primfunc_args] + outs + unbound_tir_vars
+ inputs = [*create_primfunc_args, *unbound_tir_vars, *outs]
tir_func = create_prim_func(inputs, "int64")
if primfunc_attrs:
@@ -374,8 +374,8 @@ def gen_call_tir_inputs(
for out in outs
]
- tir_vars = None
- if len(unbound_tir_vars) > 0:
- tir_vars = _shape_with_old_tir_var(unbound_tir_vars,
tir_var_inverse_map)
+ call_tir_args.extend(
+ tirx.stmt_functor.substitute(value, tir_var_inverse_map) for value in
unbound_tir_vars
+ )
- return (tir_func, call_tir_args, output_ty, tir_vars)
+ return (tir_func, call_tir_args, output_ty)
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 8c9ca99fb6..aa1c54fd51 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -401,18 +401,30 @@ class LowerTIRToLocalView : public ExprMutator {
}
std::vector<ShardingSpec> sharding_specs;
ffi::Array<Expr> args = val->args[1].as_or_throw<Tuple>()->fields;
- for (const auto& arg : args) {
- const auto* ty = GetTypeAs<DTensorTypeNode>(arg);
- TVM_FFI_ICHECK(ty);
- sharding_specs.push_back(ShardingSpec(ty->device_mesh, ty->placement));
+ GlobalVar gvar = val->args[0].as_or_throw<GlobalVar>();
+ tirx::PrimFunc prim_func = MatchPrimFunc(builder_->GetContextIRModule(),
gvar).value();
+ TVM_FFI_ICHECK_LE(args.size(), prim_func->params.size());
+ for (size_t i = 0; i < args.size(); ++i) {
+ const Expr& arg = args[i];
+ const tirx::Var& param = prim_func->params[i];
+ if (prim_func->buffer_map.count(param)) {
+ const auto* ty = GetTypeAs<DTensorTypeNode>(arg);
+ TVM_FFI_CHECK(ty, TypeError)
+ << "Expected buffer parameter " << param << " to receive a
distributed tensor, but "
+ << arg << " has type " << GetType(arg);
+ sharding_specs.push_back(ShardingSpec(ty->device_mesh, ty->placement));
+ } else {
+ TVM_FFI_CHECK(arg.as<PrimExpr>(), TypeError)
+ << "Expected scalar parameter " << param
+ << " to receive an individual primitive expression, but " << arg
<< " has type "
+ << GetType(arg);
+ }
}
Var output_var = binding->var;
ffi::Array<DTensorType> output_tys = ExtractDTensorType(output_var);
for (const auto& ty : output_tys) {
sharding_specs.push_back(ShardingSpec(ty->device_mesh, ty->placement));
}
- GlobalVar gvar = val->args[0].as_or_throw<GlobalVar>();
- tirx::PrimFunc prim_func = MatchPrimFunc(builder_->GetContextIRModule(),
gvar).value();
tirx::PrimFunc new_prim_func;
std::string allreduce_kind;
std::tie(new_prim_func, allreduce_kind) =
diff --git a/src/relax/ir/dataflow_pattern.cc b/src/relax/ir/dataflow_pattern.cc
index 7514ffceb2..5f0de5942f 100644
--- a/src/relax/ir/dataflow_pattern.cc
+++ b/src/relax/ir/dataflow_pattern.cc
@@ -650,8 +650,7 @@ ConstantPattern IsConst() { return
ConstantPattern(ffi::make_object<ConstantPatt
WildcardPattern Wildcard() { return
WildcardPattern(ffi::make_object<WildcardPatternNode>()); }
ExprPattern IsExpr(const Expr& expr) { return ExprPattern(expr); }
ExprPattern IsOp(const ffi::String& op_name) { return
IsExpr(Op::Get(op_name)); }
-CallPattern IsCallTIR(const ffi::String& name, ffi::Optional<TuplePattern>
var_args,
- ffi::Optional<DFPattern> tir_vars) {
+CallPattern IsCallTIR(const ffi::String& name, ffi::Optional<TuplePattern>
var_args) {
DFPattern arg_pattern;
if (!var_args.has_value()) {
arg_pattern = Wildcard();
@@ -659,9 +658,6 @@ CallPattern IsCallTIR(const ffi::String& name,
ffi::Optional<TuplePattern> var_a
arg_pattern = var_args.value();
}
- if (tir_vars.has_value()) {
- return IsOp("relax.call_tir")(GlobalVarPattern(name), arg_pattern,
tir_vars.value());
- }
return IsOp("relax.call_tir")(GlobalVarPattern(name), arg_pattern);
}
diff --git a/src/relax/op/distributed/distributed.cc
b/src/relax/op/distributed/distributed.cc
index 51aa1606b1..e2969517e3 100644
--- a/src/relax/op/distributed/distributed.cc
+++ b/src/relax/op/distributed/distributed.cc
@@ -110,17 +110,13 @@ Type InferTypeCallTIRLocalView(const Call& call, const
BlockBuilder& ctx) {
}
TVM_REGISTER_OP("relax.dist.call_tir_local_view")
- .set_num_inputs(3)
+ .set_num_inputs(2)
.add_argument("func", "Expr", "The destination-passing-style function.")
.add_argument("args", "Tuple", "The input arguments.")
- .add_argument("packed_ints", "Expr",
- "ShapeExpr representing a tuple of ints to unpack during
runtime. Omitted from "
- "args if unused")
.set_attr<FInferType>("FInferType", InferTypeCallTIRLocalView)
.set_attr<bool>("FPurity", true);
-Expr MakeCallTIRLocalView(Expr func, Tuple args,
ffi::Array<distributed::DTensorType> out_ty_list,
- ffi::Optional<Expr> packed_ints) {
+Expr MakeCallTIRLocalView(Expr func, Tuple args,
ffi::Array<distributed::DTensorType> out_ty_list) {
for (const distributed::DTensorType& ty : out_ty_list) {
const auto* shape = ty->tensor_ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(shape != nullptr)
@@ -137,14 +133,7 @@ Expr MakeCallTIRLocalView(Expr func, Tuple args,
ffi::Array<distributed::DTensor
}
static const Op& op = Op::Get("relax.dist.call_tir_local_view");
- Call call;
- if (!packed_ints) {
- // don't use additional optional argument
- call = Call(Type::Missing(), op, {func, args}, {}, {out_ty});
- } else {
- call = Call(Type::Missing(), op, {func, args, packed_ints.value()}, {},
{out_ty});
- }
- return call;
+ return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 08a11ef2d6..80de3c8f8e 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -274,8 +274,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
*
* \param func_ty The Type of the TIR callee.
* \param arg_ty The Type of the argument tuple.
- * \param packed_ints_ty The Type of the ffi::Shape argument,
- * if present.
* \param opt_inplace_indices For `R.call_tir_inplace`, an array of
* indices indicating which outputs are constructed from in-place
* mutation of the inputs. See
@@ -285,8 +283,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
* Otherwise, std::nullopt.
*/
static ffi::Optional<Type> InferCallTIROutputTypeFromArguments(
- Type func_ty, Type arg_ty, ffi::Optional<Type> packed_ints_ty,
- ffi::Optional<ffi::Array<int64_t>> opt_inplace_indices) {
+ Type func_ty, Type arg_ty, ffi::Optional<ffi::Array<int64_t>>
opt_inplace_indices) {
auto opt_callee_ty = func_ty.as<FuncType>();
TVM_FFI_CHECK(opt_callee_ty, TypeError)
<< "The first argument to `R.call_tir` must be a function, "
@@ -303,34 +300,18 @@ static ffi::Optional<Type>
InferCallTIROutputTypeFromArguments(
TVM_FFI_CHECK(args, TypeError) << "The second argument to `R.call_tir` must
be a tuple, "
<< "but instead received expression of type "
<< arg_ty;
- // R.call_tir expects the PrimFunc to have three groups of arguments.
+ // R.call_tir expects the PrimFunc to have two groups of arguments.
//
// 1. Input arguments that are explicitly provided as Relax arguments.
// 2. Output tensor arguments.
- // 3. Shape arguments, represented as `T.int64` in the PrimFunc, and
- // as an optional ShapeExpr argument in the `relax::Call` node.
//
// In order to determine the return type of `R.call_tir`, we must
// identify the PrimFunc arguments that will be in group (2).
size_t num_input_arguments = args->fields.size();
- size_t num_trailing_int_arguments = 0;
- const ShapeTypeNode* packed_tuple_ty = nullptr;
- if (packed_ints_ty) {
- auto packed_ty = packed_ints_ty.value();
- packed_tuple_ty = packed_ty.as<ShapeTypeNode>();
- TVM_FFI_CHECK(packed_tuple_ty && !packed_tuple_ty->IsUnknownNdim(),
TypeError)
- << "The third argument to `R.call_tir`, if present, "
- << "must be a ffi::Shape with known dimensionality. "
- << "However, the argument received was of type " << packed_ty;
- num_trailing_int_arguments = packed_tuple_ty->ndim;
- } else {
- num_trailing_int_arguments = 0;
- }
- TVM_FFI_CHECK_LE(num_input_arguments + num_trailing_int_arguments,
callee_params.size(),
- ValueError)
- << "R.call_tir attempted to call a function using " <<
num_input_arguments
- << " input arguments and " << num_trailing_int_arguments << " trailing
integer arguments. "
+ TVM_FFI_CHECK_LE(args->fields.size(), callee_params.size(), ValueError)
+ << "R.call_tir attempted to call a function using " <<
args->fields.size()
+ << " explicit arguments. "
<< "However, the callee only accepts " << callee_params.size() << "
arguments in total.";
// While Relax can specify a distributed tensor, TIR cannot. The
@@ -366,14 +347,7 @@ static ffi::Optional<Type>
InferCallTIROutputTypeFromArguments(
auto dummy_callee_ty = [&]() -> FuncType {
ffi::Array<Type> dummy_params(callee_params.begin(),
callee_params.begin() + num_input_arguments);
-
- for (size_t i = callee_params.size() - num_trailing_int_arguments; i <
callee_params.size();
- i++) {
- dummy_params.push_back(callee_params[i]);
- }
-
- ffi::Array<Type> dummy_ret(callee_params.begin() + num_input_arguments,
- callee_params.end() -
num_trailing_int_arguments);
+ ffi::Array<Type> dummy_ret(callee_params.begin() + num_input_arguments,
callee_params.end());
if (opt_inplace_indices) {
// For R.call_tir_inplace, the `inplace_indices` are used to
@@ -401,24 +375,8 @@ static ffi::Optional<Type>
InferCallTIROutputTypeFromArguments(
return FuncType(dummy_params, dummy_out_ty);
}();
- auto dummy_args = [&]() -> ffi::Array<Expr> {
- ffi::Array<Expr> dummy_args =
- args->fields.Map([](const Type& ty) -> Expr { return
Var("dummy_leading_arg", ty); });
-
- for (size_t i = 0; i < num_trailing_int_arguments; i++) {
- TVM_FFI_ICHECK(packed_tuple_ty);
- PrimType dummy_arg_ty = [&]() {
- if (packed_tuple_ty->values) {
- return PrimType(packed_tuple_ty->values.value()[i].ty());
- } else {
- return PrimType::Int(64);
- }
- }();
- dummy_args.push_back(Var("dummy_trailing_arg", dummy_arg_ty));
- }
-
- return dummy_args;
- }();
+ ffi::Array<Expr> dummy_args =
+ args->fields.Map([](const Type& ty) -> Expr { return Var("dummy_arg",
ty); });
Type derived_ret_ty = DeriveCallRetType(
dummy_callee_ty, Call(Type::Missing(), Var("dummy_callee",
dummy_callee_ty), dummy_args),
@@ -450,9 +408,8 @@ Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) {
// `relax.call_tir_inplace`. Therefore, all error messages should
// be written in terms of `call->op`, and should not explicitly
// reference the `relax.call_tir` operator.`
- TVM_FFI_ICHECK(call->args.size() == 2 || call->args.size() == 3)
- << "Operation " << call->op << " expects either two arguments [callee,
arg_tuple], "
- << "or three arguments [callee, arg_tuple, tir_args], "
+ TVM_FFI_ICHECK_EQ(call->args.size(), 2)
+ << "Operation " << call->op << " expects two arguments [callee,
arg_tuple], "
<< "but " << call << " has " << call->args.size() << " arguments.";
auto callee = call->args[0];
@@ -472,14 +429,6 @@ Expr NormalizeCallTIR(const BlockBuilder& ctx, Call call) {
<< ", which is neither an in-line tuple, "
<< "nor a variable binding that may be normalized to an in-line tuple.";
- if (call->args.size() > 2) {
- Expr packed_ints = call->args[2];
- TVM_FFI_ICHECK(packed_ints->ty.as<ShapeTypeNode>())
- << "Operation " << call->op << " expects the optional third argument, "
- << "if present, to be a ffi::Shape. "
- << "However, the third argument " << packed_ints << " has type " <<
packed_ints->ty;
- }
-
TVM_FFI_ICHECK_EQ(call->ty_args.size(), 1)
<< "R.call_tir should have exactly one `ty_args` parameter, "
<< "which defines the output of the PrimFunc.";
@@ -542,14 +491,6 @@ void ValidateCallTIR(Call call) {
auto callee = call->args[0];
Expr arg_tuple = call->args[1];
- auto packed_int_ty = [&]() -> ffi::Optional<Type> {
- if (call->args.size() <= 2) {
- return std::nullopt;
- } else {
- return GetType(call->args[2]);
- }
- }();
-
auto opt_inplace_indices = [&]() -> ffi::Optional<ffi::Array<int64_t>> {
if (const auto* attrs = call->attrs.as<CallTIRInplaceAttrs>()) {
return attrs->inplace_indices;
@@ -559,8 +500,8 @@ void ValidateCallTIR(Call call) {
}();
Type explicit_ty = call->ty_args[0];
- auto inferred_ty = InferCallTIROutputTypeFromArguments(GetType(callee),
GetType(arg_tuple),
- packed_int_ty,
opt_inplace_indices);
+ auto inferred_ty =
+ InferCallTIROutputTypeFromArguments(GetType(callee), GetType(arg_tuple),
opt_inplace_indices);
if (inferred_ty.has_value()) {
TVM_FFI_CHECK(IsBaseOf(inferred_ty.value(), explicit_ty), TypeError)
<< "The `out_ty` argument for R.call_tir must be compatible with the
PrimFunc. "
@@ -570,19 +511,15 @@ void ValidateCallTIR(Call call) {
}
TVM_REGISTER_OP("relax.call_tir")
- .set_num_inputs(3)
+ .set_num_inputs(2)
.add_argument("func", "Expr", "The destination-passing-style function.")
.add_argument("args", "Tuple", "The input arguments.")
- .add_argument("packed_ints", "Expr",
- "ShapeExpr representing a tuple of ints to unpack during
runtime. Omitted from "
- "args if unused")
.set_attr<FInferType>("FInferType", InferTypeCallTIR)
.set_attr<FNormalize>("FNormalize", NormalizeCallTIR)
.set_attr<FValidate>("FValidate", ValidateCallTIR)
.set_attr<bool>("FPurity", true);
-Expr MakeCallTIR(Expr func, Tuple args, ffi::Array<TensorType> out_ty_list,
- ffi::Optional<Expr> packed_ints) {
+Expr MakeCallTIR(Expr func, Tuple args, ffi::Array<TensorType> out_ty_list) {
for (const TensorType& ty : out_ty_list) {
const auto* shape = ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(shape != nullptr)
@@ -599,14 +536,7 @@ Expr MakeCallTIR(Expr func, Tuple args,
ffi::Array<TensorType> out_ty_list,
}
static const Op& op = Op::Get("relax.call_tir");
- Call call;
- if (!packed_ints) {
- // don't use additional optional argument
- call = Call(Type::Missing(), op, {func, args}, {}, {out_ty});
- } else {
- call = Call(Type::Missing(), op, {func, args, packed_ints.value()}, {},
{out_ty});
- }
- return call;
+ return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -617,21 +547,17 @@ TVM_FFI_STATIC_INIT_BLOCK() {
// call_tir_with_grad
TVM_REGISTER_OP("relax.call_tir_with_grad")
- .set_num_inputs(3)
+ .set_num_inputs(2)
.set_attrs_type<CallTIRWithGradAttrs>()
.add_argument("func", "Expr", "The destination-passing-style function.")
.add_argument("args", "Tuple", "The input arguments.")
- .add_argument("packed_ints", "Expr",
- "ShapeExpr representing a tuple of ints to unpack during
runtime. Omitted from "
- "args if unused")
.set_attr<FInferType>("FInferType", InferTypeCallTIR)
.set_attr<FNormalize>("FNormalize", NormalizeCallTIR)
.set_attr<FValidate>("FValidate", ValidateCallTIR)
.set_attr<bool>("FPurity", true);
Expr MakeCallTIRWithGrad(Expr func, Tuple args, ffi::Array<TensorType>
out_ty_list,
- ffi::String te_grad_name, ffi::Map<ffi::String,
ffi::Any> te_grad_kwargs,
- ffi::Optional<Expr> packed_ints) {
+ ffi::String te_grad_name, ffi::Map<ffi::String,
ffi::Any> te_grad_kwargs) {
for (const TensorType& ty : out_ty_list) {
const auto* shape = ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(shape != nullptr)
@@ -652,14 +578,7 @@ Expr MakeCallTIRWithGrad(Expr func, Tuple args,
ffi::Array<TensorType> out_ty_li
attrs->te_grad_kwargs = te_grad_kwargs;
static const Op& op = Op::Get("relax.call_tir_with_grad");
- Call call;
- if (!packed_ints) {
- // don't use additional optional argument
- call = Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
- } else {
- call = Call(Type::Missing(), op, {func, args, packed_ints.value()},
Attrs(attrs), {out_ty});
- }
- return call;
+ return Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -757,13 +676,10 @@ Expr NormalizeCallTIRInPlace(const BlockBuilder& ctx,
Call call) {
}
TVM_REGISTER_OP("relax.call_tir_inplace")
- .set_num_inputs(3)
+ .set_num_inputs(2)
.set_attrs_type<CallTIRInplaceAttrs>()
.add_argument("func", "Expr", "The destination-passing-style function.")
.add_argument("args", "Tuple", "The input arguments.")
- .add_argument("packed_ints", "Expr",
- "ShapeExpr representing a tuple of ints to unpack during
runtime. Omitted from "
- "args if unused")
.set_attr<FInferType>("FInferType", InferTypeCallTIR)
.set_attr<FNormalize>("FNormalize", NormalizeCallTIRInPlace)
.set_attr<FValidate>("FValidate", ValidateCallTIR)
@@ -773,7 +689,7 @@ TVM_REGISTER_OP("relax.call_tir_inplace")
.set_attr<bool>("FPurity", true);
Expr MakeCallTIRInplace(Expr func, Tuple args, ffi::Array<int64_t>
inplace_indices,
- ffi::Array<TensorType> out_ty_list,
ffi::Optional<Expr> packed_ints) {
+ ffi::Array<TensorType> out_ty_list) {
for (const TensorType& ty : out_ty_list) {
const auto* shape = ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(shape != nullptr)
@@ -793,14 +709,7 @@ Expr MakeCallTIRInplace(Expr func, Tuple args,
ffi::Array<int64_t> inplace_indic
}
static const Op& op = Op::Get("relax.call_tir_inplace");
- Call call;
- if (!packed_ints) {
- // don't use additional optional argument
- call = Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
- } else {
- call = Call(Type::Missing(), op, {func, args, packed_ints.value()},
Attrs(attrs), {out_ty});
- }
- return call;
+ return Call(Type::Missing(), op, {func, args}, Attrs(attrs), {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/script/builder/distributed.cc
b/src/relax/script/builder/distributed.cc
index 496b33606b..dc9440bdee 100644
--- a/src/relax/script/builder/distributed.cc
+++ b/src/relax/script/builder/distributed.cc
@@ -28,8 +28,7 @@
namespace tvm {
namespace relax {
-Expr MakeCallTIRDist(Expr func, Tuple args,
ffi::Array<distributed::DTensorType> out_ty_list,
- ffi::Optional<Expr> packed_ints) {
+Expr MakeCallTIRDist(Expr func, Tuple args,
ffi::Array<distributed::DTensorType> out_ty_list) {
for (const distributed::DTensorType& ty : out_ty_list) {
const auto* shape = ty->tensor_ty->shape.as<ShapeExprNode>();
TVM_FFI_ICHECK(shape != nullptr)
@@ -46,14 +45,7 @@ Expr MakeCallTIRDist(Expr func, Tuple args,
ffi::Array<distributed::DTensorType>
}
static const Op& op = Op::Get("relax.call_tir");
- Call call;
- if (!packed_ints) {
- // don't use additional optional argument
- call = Call(Type::Missing(), op, {func, args}, {}, {out_ty});
- } else {
- call = Call(Type::Missing(), op, {func, args, packed_ints.value()}, {},
{out_ty});
- }
- return call;
+ return Call(Type::Missing(), op, {func, args}, {}, {out_ty});
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index 9b3b41e565..b330fd3c03 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -82,7 +82,7 @@ ffi::Optional<ExprDoc> PrintCallTIRDPSPacked(const Call& n,
const AccessPath& n_
!n->op.same_as(call_tir_inplace_op)) {
return std::nullopt;
}
- TVM_FFI_ICHECK(n->args.size() == 2 || n->args.size() == 3);
+ TVM_FFI_ICHECK_EQ(n->args.size(), 2);
TVM_FFI_ICHECK(n->ty_args.size() == 1);
ffi::Array<ExprDoc> args;
ffi::Array<ffi::String> kwargs_keys;
@@ -145,11 +145,6 @@ ffi::Optional<ExprDoc> PrintCallTIRDPSPacked(const Call&
n, const AccessPath& n_
if (n->op.same_as(call_dps_packed_op)) {
return Relax(d, "call_dps_packed")->Call(args, kwargs_keys, kwargs_values);
}
- // Step 4. Print n->args[2], the tirx variables
- if (n->args.size() == 3) {
- kwargs_keys.push_back("tir_vars");
- kwargs_values.push_back(d->AsDoc<ExprDoc>(n->args[2],
n_p->Attr("args")->ArrayItem(2)));
- }
if (n->op.same_as(call_tir_local_view)) {
return Relax(d, "dist.call_tir_local_view")->Call(args, kwargs_keys,
kwargs_values);
} else if (is_dtensor) {
diff --git a/src/relax/transform/call_tir_rewrite.cc
b/src/relax/transform/call_tir_rewrite.cc
index 8c1a16c66a..e1d9a265dc 100644
--- a/src/relax/transform/call_tir_rewrite.cc
+++ b/src/relax/transform/call_tir_rewrite.cc
@@ -68,8 +68,6 @@ class CallTIRMutator : public ExprMutator {
static const Op& call_tir_inplace_op = Op::Get("relax.call_tir_inplace");
static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
static const Op& alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor");
- static const Op& call_tir_dyn_op = Op::Get("relax.vm.call_tir_dyn");
-
if (call->op.same_as(call_tir_op) || call->op.same_as(call_tir_inplace_op)
||
call->op.same_as(call_dps_packed_op)) {
bool is_inplace = call->op.same_as(call_tir_inplace_op);
@@ -157,14 +155,7 @@ class CallTIRMutator : public ExprMutator {
}
}
}
-
- if (call->args.size() == 2) {
- builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
- } else {
- // unpack semantics
- args.push_back(call->args[2]);
- builder_->Emit(Call(Type::Missing(), call_tir_dyn_op,
{call->args[0], Tuple(args)}), "_");
- }
+ builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
} else {
if (!is_inplace) {
args = outs;
diff --git a/src/relax/transform/fold_constant.cc
b/src/relax/transform/fold_constant.cc
index 3930cf33a3..274ad7b0a5 100644
--- a/src/relax/transform/fold_constant.cc
+++ b/src/relax/transform/fold_constant.cc
@@ -280,9 +280,6 @@ class ConstantFolder : public ExprMutator {
if (!func || !arr_args) return {};
- // tir_vars are passed as extra scalar arguments to the PrimFunc, which we
cannot supply here.
- if (call->args.size() > 2) return {};
-
// Handle tuple output: ty_args[0] is a TupleType.
if (const auto* tuple_ty = call->ty_args[0].as<TupleTypeNode>()) {
return ConstEvaluateCallTIRTuple(func.value(), arr_args.value(),
tuple_ty);
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 54f7e5af35..e74584856d 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -578,16 +578,9 @@ class FusedTIRConstructor : public ExprVisitor {
func_info_.expr2buffers.Set(relax_param, param_buffers);
}
- // Move all scalar params after buffer params. To ensure that the
- // order is deterministic and predictable for testing purposes,
- // std::stable_sort is used instead of std::sort.
- std::stable_sort(prim_func_params.begin(), prim_func_params.end(),
- [](const auto& a, const auto& b) {
- bool a_is_var = a.template
as<tirx::PrimVar>().has_value();
- bool b_is_var = b.template
as<tirx::PrimVar>().has_value();
- return a_is_var < b_is_var;
- });
-
+ // Preserve the Relax function's parameter order. Tensor and primitive
+ // parameters are both explicit call_tir arguments, while output buffers
+ // are appended after the complete explicit argument prefix.
for (const auto& param : prim_func_params) {
if (auto opt = param.as<tirx::Buffer>()) {
auto buffer = opt.value();
@@ -600,6 +593,8 @@ class FusedTIRConstructor : public ExprVisitor {
tirx::Var param = tirx::Var("p_" + buffer->name,
PointerType::VoidPointerTy());
func_info_.params.push_back(param);
func_info_.buffer_map.Set(param, buffer);
+ } else if (auto var = param.as<tirx::PrimVar>()) {
+ func_info_.params.push_back(var.value());
}
}
@@ -649,14 +644,7 @@ class FusedTIRConstructor : public ExprVisitor {
func_info_.output_buffers.insert(buffers[i].get());
}
- // Step 4. Append symbolic vars
- for (const auto& param : prim_func_params) {
- if (auto var = param.as<tirx::PrimVar>()) {
- func_info_.params.push_back(var.value());
- }
- }
-
- // Step 5. Create PrimFunc
+ // Step 4. Create PrimFunc
fused_tir_ = ConstructFunc();
}
@@ -711,23 +699,6 @@ class FusedTIRConstructor : public ExprVisitor {
AllocateIntermediateBuffer(call, prim_func, output_buffer_shapes);
- // Step 6. Update tir_vars
- if (call->args.size() > 2) {
- TVM_FFI_ICHECK(call->args.size() == 3);
- const Expr& tir_vars = call->args[2];
- if (const auto* shape_expr = tir_vars.as<ShapeExprNode>()) {
- const auto& args = shape_expr->values;
- size_t num_params = prim_func->params.size();
- TVM_FFI_ICHECK_GE(num_params, args.size());
- for (size_t i = 0; i < args.size(); ++i) {
- const tirx::Var& param = prim_func->params[num_params - args.size()
+ i];
- func_info_.symbolic_var_matcher.Match(param.as_or_throw<PrimExpr>(),
args[i]);
- }
- } else {
- TVM_FFI_THROW(InternalError)
- << "TIR vars should be a shape expr, but got: " <<
tir_vars->GetTypeKey();
- }
- }
// Update fused func name
func_info_.global_name += "_" + gv->name_hint;
}
@@ -833,17 +804,23 @@ class FusedTIRConstructor : public ExprVisitor {
void MapInputBuffer(const tirx::PrimFunc& func, const relax::Expr& args) {
ffi::Array<Expr> arg_list;
ffi::Array<tirx::Buffer> buffer_list;
- if (const auto* arg_tuple = args.as<TupleNode>()) {
- arg_list = arg_tuple->fields;
- } else {
- arg_list = {args};
- }
+ ffi::Array<Expr> call_args = args.as_or_throw<Tuple>()->fields;
- TVM_FFI_ICHECK_GE(func->params.size(), arg_list.size());
- for (size_t i = 0; i < arg_list.size(); ++i) {
+ TVM_FFI_ICHECK_GE(func->params.size(), call_args.size());
+ for (size_t i = 0; i < call_args.size(); ++i) {
+ const Expr& arg = call_args[i];
const tirx::Var& param = func->params[i];
- const tirx::Buffer& buffer = func->buffer_map.at(param);
- buffer_list.push_back(buffer);
+ if (func->buffer_map.count(param)) {
+ arg_list.push_back(arg);
+ buffer_list.push_back(func->buffer_map.at(param));
+ } else {
+ auto prim_arg = arg.as<PrimExpr>();
+ TVM_FFI_CHECK(prim_arg.has_value(), TypeError)
+ << "Expected scalar parameter " << param
+ << " to receive an individual primitive expression, but " << arg
<< " has type "
+ << GetType(arg);
+ func_info_.symbolic_var_matcher.Match(param.as_or_throw<PrimExpr>(),
prim_arg.value());
+ }
}
MapArgsToBuffer(arg_list, buffer_list);
@@ -852,7 +829,6 @@ class FusedTIRConstructor : public ExprVisitor {
static ffi::Array<tirx::Var> GetPrimFuncOutputParams(const tirx::PrimFunc&
func,
const
ffi::Array<int64_t>& output_indices) {
size_t n = func->params.size();
- int symbolic_var_index = -1;
size_t output_size = output_indices.size();
TVM_FFI_ICHECK_GE(n, output_size);
@@ -860,24 +836,11 @@ class FusedTIRConstructor : public ExprVisitor {
for (int64_t idx : output_indices) {
int i = static_cast<int>(idx);
const tirx::Var& param = func->params[static_cast<size_t>(i)];
- auto param_ty = param->ty.as<PrimType>();
- if (param_ty && (param_ty.value().code() == DLDataTypeCode::kDLInt ||
- param_ty.value().code() == DLDataTypeCode::kDLUInt)) {
- if (symbolic_var_index == -1) symbolic_var_index = i;
- } else if (param->ty.as<PointerTypeNode>()) {
- TVM_FFI_ICHECK(symbolic_var_index == -1)
- << "The scalar input should be at the ending of the "
- "parameter list.";
- ret.push_back(param);
- } else {
- TVM_FFI_THROW(InternalError)
- << "The params of PrimFunc are expected to be Buffer handle or
scalar, but got: "
- << param->ty;
- }
+ TVM_FFI_ICHECK(param->ty.as<PointerTypeNode>())
+ << "The output params of a PrimFunc must be buffer handles, but
parameter " << i
+ << " has type " << param->ty;
+ ret.push_back(param);
}
-
- size_t end_index = symbolic_var_index == -1 ? n : symbolic_var_index;
- TVM_FFI_ICHECK_GE(end_index, output_size);
return ret;
}
@@ -909,16 +872,16 @@ class FusedTIRConstructor : public ExprVisitor {
}
ffi::Array<tirx::Var> output_params = GetPrimFuncOutputParams(func,
output_idxs);
- auto input_buffers = func_info_.expr2buffers.Get(call->args[1]);
for (size_t i = 0; i < output_size; ++i) {
const tirx::Var& param = output_params[i];
const tirx::Buffer& buffer = func->buffer_map.at(param);
// if this is an inplace output, do not do an intermediate allocation
if (output_idxs[i] < num_inputs) {
- TVM_FFI_ICHECK(input_buffers.has_value())
- << "Inplace functions must have some defined input";
- output_buffers.push_back(input_buffers.value()[output_idxs[i]]);
+ auto it = func_info_.buffer_subst_map.find(buffer);
+ TVM_FFI_ICHECK(it != func_info_.buffer_subst_map.end())
+ << "Inplace output buffer " << buffer << " must be mapped to a
defined input";
+ output_buffers.push_back((*it).second);
continue;
}
@@ -1241,7 +1204,6 @@ class TIRFuseMutator : public ExprMutator {
// are not supported by PrimFunc, so this step verifies that
// ExpandTupleArguments has already removed them.
ffi::Array<Expr> arg_list;
- ffi::Array<PrimExpr> tir_vars;
for (size_t i = 0; i < call->args.size(); ++i) {
auto arg = call->args[i];
auto ty = GetType(arg);
@@ -1260,11 +1222,11 @@ class TIRFuseMutator : public ExprMutator {
for (const PrimExpr& prim_value : shape->values.value()) {
TVM_FFI_ICHECK(prim_value.as<tirx::PrimVar>())
<< "All shape inputs are expected to be single tirx var.";
- tir_vars.push_back(prim_value);
+ arg_list.push_back(prim_value);
}
} else if (ty.as<PrimTypeNode>()) {
if (auto literal = arg.as<PrimExpr>()) {
- tir_vars.push_back(literal.value());
+ arg_list.push_back(literal.value());
} else {
TVM_FFI_THROW(TypeError) << "FuseTIR expects scalar arguments to be
PrimExpr, "
<< "but received " << arg;
@@ -1277,9 +1239,6 @@ class TIRFuseMutator : public ExprMutator {
// Step b. Create call_tir or call_tir_inplace
ffi::Array<Expr> call_args = {fused_tir_gv, Tuple(arg_list)};
- if (!tir_vars.empty()) {
- call_args.push_back(ShapeExpr(tir_vars));
- }
Op call_op = call_tir_op_;
Attrs call_attrs = call->attrs;
if (replacement.inplace_indices.size()) {
diff --git
a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
index a7b9061b57..1cf1f0bb92 100644
---
a/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
+++
b/tests/python/relax/distributed/test_distributed_transform_lower_global_to_local_view.py
@@ -715,13 +715,11 @@ def test_llama_attention():
cls.rotary_embedding,
(lv9, cos_cached, sin_cached),
out_ty=R.DTensor((1, 256, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([256]),
)
lv17 = R.dist.call_tir(
cls.rotary_embedding,
(lv12, cos_cached, sin_cached),
out_ty=R.DTensor((1, 256, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([256]),
)
lv18 = R.dist.call_tir(
cls.reshape1,
@@ -1406,7 +1404,6 @@ def test_llama_attention():
cls.rotary_embedding1,
(lv9, cos_cached, sin_cached),
out_ty=R.DTensor((1, 256, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([256]),
)
)
lv17: R.DTensor((1, 256, 32, 128), "float16", "mesh[0]", "S[2]") =
(
@@ -1414,7 +1411,6 @@ def test_llama_attention():
cls.rotary_embedding1,
(lv12, cos_cached, sin_cached),
out_ty=R.DTensor((1, 256, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([256]),
)
)
lv18: R.DTensor((256, 32, 128), "float16", "mesh[0]", "S[1]") = (
diff --git
a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
index 7f8b1cadcf..bcfe0d34c7 100644
---
a/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
+++
b/tests/python/relax/distributed/test_distributed_transform_propagate_sharding.py
@@ -1595,8 +1595,8 @@ def test_decoder_layer_dynamic_shape():
var_A: T.handle,
B: T.Buffer((T.int64(2048), T.int64(128)), "float16"),
C: T.Buffer((T.int64(2048), T.int64(128)), "float16"),
- var_rotary: T.handle,
m: T.int64,
+ var_rotary: T.handle,
):
T.func_attr({"tirx.noalias": True})
n = T.int64()
@@ -1673,15 +1673,13 @@ def test_decoder_layer_dynamic_shape():
)
lv16 = R.call_tir(
cls.rotary_embedding,
- (lv9, cos_cached, sin_cached),
+ (lv9, cos_cached, sin_cached, m),
out_ty=R.Tensor((1, n, 32, 128), dtype="float16"),
- tir_vars=R.shape([m]),
)
lv17 = R.call_tir(
cls.rotary_embedding,
- (lv12, cos_cached, sin_cached),
+ (lv12, cos_cached, sin_cached, m),
out_ty=R.Tensor((1, n, 32, 128), dtype="float16"),
- tir_vars=R.shape([m]),
)
lv18: R.Tensor((n, 32, 128), dtype="float16") = R.reshape(lv17,
R.shape([n, 32, 128]))
lv19: R.Tensor((n, 32, 128), dtype="float16") = R.reshape(lv15,
R.shape([n, 32, 128]))
@@ -1797,8 +1795,8 @@ def test_decoder_layer_dynamic_shape():
var_A: T.handle,
B: T.Buffer((T.int64(2048), T.int64(128)), "float16"),
C: T.Buffer((T.int64(2048), T.int64(128)), "float16"),
- var_rotary: T.handle,
m: T.int64,
+ var_rotary: T.handle,
):
T.func_attr({"tirx.noalias": True})
n = T.int64()
@@ -1868,15 +1866,13 @@ def test_decoder_layer_dynamic_shape():
)
lv16 = R.dist.call_tir(
cls.rotary_embedding,
- (lv9, cos_cached, sin_cached),
+ (lv9, cos_cached, sin_cached, m),
out_ty=R.DTensor((1, n, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([m]),
)
lv17 = R.dist.call_tir(
cls.rotary_embedding,
- (lv12, cos_cached, sin_cached),
+ (lv12, cos_cached, sin_cached, m),
out_ty=R.DTensor((1, n, 32, 128), "float16", "mesh[0]",
"S[2]"),
- tir_vars=R.shape([m]),
)
lv18: R.DTensor((n, 32, 128), "float16", "mesh[0]", "S[1]") =
R.reshape(
lv17, R.shape([n, 32, 128])
diff --git a/tests/python/relax/test_analysis_well_formed.py
b/tests/python/relax/test_analysis_well_formed.py
index da81d642aa..ed37b367e2 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -748,6 +748,76 @@ def test_call_tir_with_matching_arguments():
rx.analysis.well_formed(Module)
+def test_call_tir_with_interspersed_primitive_argument():
+ """Primitive values are positional arguments in the call_tir tuple."""
+
+ @I.ir_module(s_tir=True)
+ class Module:
+ @R.function
+ def main(
+ A: R.Tensor([16], "float16"),
+ scale: R.Prim("float32"),
+ C: R.Tensor([16], "float16"),
+ ):
+ B = R.call_tir(Module.add_scaled, (A, scale, C),
out_ty=R.Tensor([16], "float16"))
+ return B
+
+ @T.prim_func(s_tir=True)
+ def add_scaled(
+ A: T.Buffer([T.int64(16)], "float16"),
+ scale: T.float32,
+ C: T.Buffer([T.int64(16)], "float16"),
+ B: T.Buffer([T.int64(16)], "float16"),
+ ):
+ for i in range(T.int64(16)):
+ B[i] = A[i] + T.Cast("float16", scale) * C[i]
+
+ rx.analysis.well_formed(Module)
+
+
+def test_call_tir_with_incorrect_primitive_argument_dtype():
+ """Primitive call_tir arguments must match the PrimFunc parameter dtype."""
+
+ @I.ir_module(check_well_formed=False, s_tir=True)
+ class Module:
+ @R.function
+ def main(A: R.Tensor([16], "float16"), scale: R.Prim("int64")):
+ B = R.call_tir(Module.scale, (A, scale), out_ty=R.Tensor([16],
"float16"))
+ return B
+
+ @T.prim_func(s_tir=True)
+ def scale(
+ A: T.Buffer([T.int64(16)], "float16"),
+ scale: T.float32,
+ B: T.Buffer([T.int64(16)], "float16"),
+ ):
+ for i in range(T.int64(16)):
+ B[i] = A[i] * T.Cast("float16", scale)
+
+ assert not rx.analysis.check_well_formed(Module)
+
+
+def test_call_tir_shape_expr_is_not_a_primitive_argument():
+ """ShapeExpr groups must be unpacked into positional primitive values."""
+
+ @I.ir_module(check_well_formed=False, s_tir=True)
+ class Module:
+ @R.function
+ def main():
+ B = R.call_tir(
+ Module.make_tensor,
+ (R.shape([1, 2]),),
+ out_ty=R.Tensor([1], "float32"),
+ )
+ return B
+
+ @T.prim_func(s_tir=True)
+ def make_tensor(m: T.int64, n: T.int64, B: T.Buffer([T.int64(1)],
"float32")):
+ B[0] = T.Cast("float32", m + n)
+
+ assert not rx.analysis.check_well_formed(Module)
+
+
def test_call_tir_input_ndim():
"""Arguments to R.call_tir must have the correct dimensionality
diff --git a/tests/python/relax/test_blockbuilder_emit_te.py
b/tests/python/relax/test_blockbuilder_emit_te.py
index 6c7f192373..b7bb6eec94 100644
--- a/tests/python/relax/test_blockbuilder_emit_te.py
+++ b/tests/python/relax/test_blockbuilder_emit_te.py
@@ -47,8 +47,8 @@ def test_emit_te_with_symbolic_arg():
@T.prim_func(private=True, s_tir=True)
def te_func(
A: T.Buffer((T.int64(10),), "float32"),
- B: T.Buffer((T.int64(10),), "float32"),
m: T.int64,
+ B: T.Buffer((T.int64(10),), "float32"),
):
T.func_attr({"tirx.noalias": True})
for i in range(T.int64(10)):
@@ -65,9 +65,8 @@ def test_emit_te_with_symbolic_arg():
cls = Expected
gv = R.call_tir(
cls.te_func,
- (x,),
+ (x, m),
out_ty=R.Tensor((10,), dtype="float32"),
- tir_vars=R.shape([m]),
)
return gv
@@ -75,7 +74,7 @@ def test_emit_te_with_symbolic_arg():
def test_symbolic_shape_in_prim_value():
- """Scalar primitive Vars flow through the optimistic tir_vars path."""
+ """Scalar primitive Vars become ordinary call_tir arguments."""
def te_slice(tensor, i):
return tvm.te.compute([tensor.shape[1]], lambda j: tensor[i, j],
name="slice")
@@ -96,8 +95,8 @@ def test_symbolic_shape_in_prim_value():
@T.prim_func(private=True, s_tir=True)
def te_slice(
A: T.Buffer([T.int64(16), T.int64(16)], "float32"),
- Output: T.Buffer(T.int64(16), "float32"),
row_index: T.int64,
+ Output: T.Buffer(T.int64(16), "float32"),
):
T.func_attr({"tirx.noalias": True})
@@ -115,9 +114,8 @@ def test_symbolic_shape_in_prim_value():
gv = R.call_tir(
cls.te_slice,
- (A,),
+ (A, arg_row_index),
out_ty=R.Tensor([16], "float32"),
- tir_vars=R.shape([arg_row_index]),
)
return gv
diff --git a/tests/python/relax/test_dataflow_pattern.py
b/tests/python/relax/test_dataflow_pattern.py
index c507ea6f46..57b81444a7 100644
--- a/tests/python/relax/test_dataflow_pattern.py
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -59,7 +59,7 @@ class Module:
B[vi, vj] = T.max(A[vi, vj], 0.0)
@T.prim_func(s_tir=True)
- def tir_zeros(x: T.handle, n: T.int64):
+ def tir_zeros(n: T.int64, x: T.handle):
T.func_attr({"global_symbol": "tir_zeros"})
A = T.match_buffer(x, [n])
for i in range(n):
@@ -73,9 +73,7 @@ class Module:
with R.dataflow():
lv0 = R.call_tir(cls.tir_matmul, (x, w), R.Tensor((32, 32),
dtype="float32"))
lv1 = R.call_tir(cls.tir_relu, (lv0), R.Tensor((32, 32),
dtype="float32"))
- lv2 = R.call_tir(
- cls.tir_zeros, [], R.Tensor((32,), dtype="float32"),
tir_vars=R.ShapeExpr([32])
- )
+ lv2 = R.call_tir(cls.tir_zeros, [32], R.Tensor((32,),
dtype="float32"))
gv = (lv1, lv2)
R.output(gv)
return gv
@@ -305,7 +303,7 @@ def test_is_call_tir():
assert is_call_tir("tir_relu").match(lv1_val)
assert is_call_tir("tir_relu", [is_call_tir("tir_matmul")]).match(lv1_val,
var2val=var2val)
assert not is_call_tir("tir_relu",
[is_call_tir("tir_relu")]).match(lv1_val, var2val=var2val)
- assert is_call_tir("tir_zeros", wildcard(), wildcard()).match(lv2_val,
var2val=var2val)
+ assert is_call_tir("tir_zeros", [wildcard()]).match(lv2_val,
var2val=var2val)
@R.function(pure=False)
diff --git a/tests/python/relax/test_frontend_nn_op.py
b/tests/python/relax/test_frontend_nn_op.py
index 2877ef11c7..ef313b95b5 100644
--- a/tests/python/relax/test_frontend_nn_op.py
+++ b/tests/python/relax/test_frontend_nn_op.py
@@ -638,18 +638,10 @@ def test_tensor_ir_op():
@T.prim_func(private=True, s_tir=True)
def fused_rope( # pylint: disable=too-many-locals
var_qkv: T.handle,
+ offset: T.int64,
var_q: T.handle,
var_k: T.handle,
var_v: T.handle,
- # Scalar arguments must be specified after tensor arguments,
- # including the output tensor arguments
- #
- # TODO(Lunderberg): Update
- # `tvm.relax.frontend.nn.op.tensor_ir_op` to use `Expr`
- # instead of `tir_vars`, so that the order can be consistent
- # between the function definition and the arguments in
- # `op.tensor_ir_op`.
- offset: T.int64,
):
batch_size = T.int64()
seq_len = T.int64()
@@ -677,7 +669,7 @@ def test_tensor_ir_op():
@I.ir_module(s_tir=True)
class Expected:
@T.prim_func(private=True, s_tir=True)
- def llama_fused_rope(var_qkv: T.handle, var_q: T.handle, var_k:
T.handle, var_v: T.handle, offset: T.int64):
+ def llama_fused_rope(var_qkv: T.handle, offset: T.int64, var_q:
T.handle, var_k: T.handle, var_v: T.handle):
batch_size, seq_len = T.int64(), T.int64()
qkv = T.match_buffer(var_qkv, (batch_size, seq_len, 24, 16),
"float16")
q = T.match_buffer(var_q, (batch_size, seq_len, 8, 16), "float16")
@@ -700,7 +692,7 @@ def test_tensor_ir_op():
R.func_attr({"num_input": 3})
cls = Expected
with R.dataflow():
- lv1 = R.call_tir(cls.llama_fused_rope, (qkv,),
out_ty=[R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16),
dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16")],
tir_vars=R.shape([offset_1]))
+ lv1 = R.call_tir(cls.llama_fused_rope, (qkv, offset_1),
out_ty=[R.Tensor((1, 1, 8, 16), dtype="float16"), R.Tensor((1, 1, 8, 16),
dtype="float16"), R.Tensor((1, 1, 8, 16), dtype="float16")])
llama_fused_rope_0: R.Tensor((1, 1, 8, 16), dtype="float16") =
lv1[0]
llama_fused_rope_1: R.Tensor((1, 1, 8, 16), dtype="float16") =
lv1[1]
llama_fused_rope_2: R.Tensor((1, 1, 8, 16), dtype="float16") =
lv1[2]
@@ -798,10 +790,9 @@ def test_tensor_ir_inplace_op():
with R.dataflow():
lv1 = R.call_tir_inplace(
cls.inplace_take,
- (embedding_table, input_ids, embedding_dst),
+ (embedding_table, input_ids, embedding_dst, offset_1),
out_ty=R.Tensor((total_seq_len, hidden_size), dtype),
inplace_indices=[2],
- tir_vars=R.shape([offset_1]),
)
gv1: R.Tensor((total_seq_len, hidden_size), dtype) = lv1
R.output(gv1)
diff --git a/tests/python/relax/test_op_index.py
b/tests/python/relax/test_op_index.py
index c102b0fdc6..f70e542b95 100644
--- a/tests/python/relax/test_op_index.py
+++ b/tests/python/relax/test_op_index.py
@@ -911,16 +911,15 @@ def test_legalize_dynamic_begin_end():
index = T.int64()
return R.call_tir(
expected.strided_slice,
- (A,),
+ (A, index),
out_ty=R.Tensor((1, 16), "float32"),
- tir_vars=R.shape([index]),
)
@T.prim_func(private=True, s_tir=True)
def strided_slice(
A: T.Buffer((T.int64(16), T.int64(16))),
- B: T.Buffer((T.int64(1), T.int64(16))),
index: T.int64,
+ B: T.Buffer((T.int64(1), T.int64(16))),
):
T.func_attr({"tirx.noalias": True})
for iters in T.grid(*B.shape):
@@ -948,7 +947,7 @@ def test_legalize_dynamic_begin_inf_end():
@I.ir_module(s_tir=True)
class expected:
@T.prim_func(private=True, s_tir=True)
- def strided_slice(A: T.Buffer((T.int64(16), T.int64(16)), "float32"),
var_T_dynamic_strided_slice_with_axes: T.handle, index: T.int64):
+ def strided_slice(A: T.Buffer((T.int64(16), T.int64(16)), "float32"),
index: T.int64, var_T_dynamic_strided_slice_with_axes: T.handle):
T.func_attr({"tirx.noalias": True})
T_dynamic_strided_slice_with_axes =
T.match_buffer(var_T_dynamic_strided_slice_with_axes, (T.max(T.int64(16) -
T.max(T.if_then_else(index < T.int64(0), index + T.int64(16), index),
T.int64(0)), T.int64(0)), T.int64(16)))
# with T.sblock("root"):
@@ -963,7 +962,7 @@ def test_legalize_dynamic_begin_inf_end():
def main(A: R.Tensor((16, 16), dtype="float32"), B:
R.Shape(["index"])) -> R.Tensor(("T.max(16 - T.max(T.if_then_else(index < 0,
index + 16, index), 0), 0)", 16), dtype="float32"):
index = T.int64()
cls = expected
- gv = R.call_tir(cls.strided_slice, (A,), out_ty=R.Tensor((T.max(16
- T.max(T.if_then_else(index < 0, index + 16, index), 0), 0), 16),
dtype="float32"), tir_vars=R.shape([index]))
+ gv = R.call_tir(cls.strided_slice, (A, index),
out_ty=R.Tensor((T.max(16 - T.max(T.if_then_else(index < 0, index + 16, index),
0), 0), 16), dtype="float32"))
return gv
# fmt: on
diff --git a/tests/python/relax/test_transform.py
b/tests/python/relax/test_transform.py
index 746ac5d858..d32caa69d5 100644
--- a/tests/python/relax/test_transform.py
+++ b/tests/python/relax/test_transform.py
@@ -178,6 +178,42 @@ def test_call_tir_rewrite():
assert s2.op.name_hint == "exp"
+def test_call_tir_rewrite_with_interspersed_primitive_argument():
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def scale_add(
+ A: T.Buffer((16,), "float32"),
+ scale: T.float32,
+ C: T.Buffer((16,), "float32"),
+ B: T.Buffer((16,), "float32"),
+ ):
+ for i in range(16):
+ B[i] = A[i] + scale * C[i]
+
+ @R.function
+ def main(
+ A: R.Tensor((16,), "float32"),
+ scale: R.Prim("float32"),
+ C: R.Tensor((16,), "float32"),
+ ) -> R.Tensor((16,), "float32"):
+ R.func_attr({"relax.force_pure": True})
+ B = R.call_tir(Module.scale_add, (A, scale, C), R.Tensor((16,),
"float32"))
+ return B
+
+ after = relax.transform.CallTIRRewrite()(Module)
+ func = after["main"]
+ bindings = func.body.blocks[0].bindings
+ output_buffer = bindings[0].var
+ call = bindings[1].value
+
+ assert call.op.name_hint == "scale_add"
+ tvm.ir.assert_structural_equal(
+ call.args,
+ [func.params[0], func.params[1], func.params[2], output_buffer],
+ )
+
+
def test_transform_remove_purity_checking():
@tvm.script.ir_module
class Before:
diff --git a/tests/python/relax/test_transform_fold_constant.py
b/tests/python/relax/test_transform_fold_constant.py
index f6720131c0..f580924329 100644
--- a/tests/python/relax/test_transform_fold_constant.py
+++ b/tests/python/relax/test_transform_fold_constant.py
@@ -585,13 +585,13 @@ def test_fold_large_op_with_tensor_input():
tvm.ir.assert_structural_equal(after, expected)
-def test_call_tir_with_tir_vars_not_folded():
- """call_tir with symbolic tir_vars cannot be const-evaluated."""
+def test_call_tir_with_primitive_args_not_folded():
+ """call_tir with symbolic primitive arguments cannot be const-evaluated."""
@tvm.script.ir_module
class Module:
@T.prim_func(private=True, s_tir=True)
- def shape_to_tensor(out: T.Buffer((T.int64(1),), "int64"), m: T.int64):
+ def shape_to_tensor(m: T.int64, out: T.Buffer((T.int64(1),), "int64")):
for i in range(T.int64(1)):
with T.sblock("out"):
vi = T.axis.remap("S", [i])
@@ -601,9 +601,7 @@ def test_call_tir_with_tir_vars_not_folded():
def main(x: R.Tensor(("m",), "float32")):
m = T.int64()
cls = Module
- gv = relax.call_tir(
- cls.shape_to_tensor, R.tuple(), R.Tensor((1,), "int64"),
tir_vars=R.shape([m])
- )
+ gv = relax.call_tir(cls.shape_to_tensor, (m,), R.Tensor((1,),
"int64"))
return gv
after = relax.transform.FoldConstant()(Module)
diff --git a/tests/python/relax/test_transform_fuse_tir.py
b/tests/python/relax/test_transform_fuse_tir.py
index 96aa496e5e..006621e4a4 100644
--- a/tests/python/relax/test_transform_fuse_tir.py
+++ b/tests/python/relax/test_transform_fuse_tir.py
@@ -959,8 +959,8 @@ def test_symbolic_var_in_call_tir_args():
def foo(
X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)),
"float32"),
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
- rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32),
T.int64(128)), "float32"),
m: T.int64,
+ rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32),
T.int64(128)), "float32"),
):
for i0, i1, i2, i3 in T.grid(T.int64(1), T.int64(1), T.int64(32),
T.int64(128)):
with T.sblock("rotary"):
@@ -980,9 +980,8 @@ def test_symbolic_var_in_call_tir_args():
lv1 = R.emit_te(topi.add, x, x)
gv = R.call_tir(
cls.foo,
- [lv1, y],
+ [lv1, y, m],
out_ty=R.Tensor((1, 1, 32, 128), dtype="float32"),
- tir_vars=R.shape([m]),
)
R.output(gv)
return gv
@@ -1005,8 +1004,8 @@ def test_symbolic_var_in_call_tir_args():
def fused(
X: T.Buffer((T.int64(1), T.int64(1), T.int64(32), T.int64(128)),
"float32"),
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
- rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32),
T.int64(128)), "float32"),
m: T.int64,
+ rotary: T.Buffer((T.int64(1), T.int64(1), T.int64(32),
T.int64(128)), "float32"),
):
T.func_attr({"tirx.noalias": True})
T_add = T.sblock_alloc_buffer((T.int64(1), T.int64(1),
T.int64(32), T.int64(128)))
@@ -1032,9 +1031,8 @@ def test_symbolic_var_in_call_tir_args():
with R.dataflow():
gv = R.call_tir(
cls.fused,
- (x, y),
+ (x, y, m),
out_ty=R.Tensor([1, 1, 32, 128], "float32"),
- tir_vars=R.shape([m]),
)
R.output(gv)
return gv
@@ -1196,8 +1194,8 @@ def test_tir_expression_in_shape():
def fused_transpose_matmul(
x: T.Buffer((T.int64(3), T.int64(4)), "float32"),
p_y: T.handle,
- p_output0: T.handle,
n: T.int64,
+ p_output0: T.handle,
):
T.func_attr({"tirx.noalias": True})
y = T.match_buffer(p_y, (n - T.int64(1), T.int64(4)))
@@ -1228,9 +1226,8 @@ def test_tir_expression_in_shape():
with R.dataflow():
lv = R.call_tir(
cls.fused_transpose_matmul,
- (x, y),
+ (x, y, n),
out_ty=R.Tensor((n - 1, 3), dtype="float32"),
- tir_vars=R.shape([n]),
)
R.output(lv)
return lv
@@ -1466,8 +1463,8 @@ def test_symbolic_var_in_buffer_shape():
def foo(
X_handle: T.handle,
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
- rotary_handle: T.handle,
m: T.int64,
+ rotary_handle: T.handle,
):
sequence_length = T.int64()
@@ -1497,9 +1494,8 @@ def test_symbolic_var_in_buffer_shape():
lv1 = R.emit_te(topi.add, x, x)
gv = R.call_tir(
cls.foo,
- [lv1, y],
+ [lv1, y, m],
out_ty=R.Tensor((1, sequence_length, 32, 128),
dtype="float32"),
- tir_vars=R.shape([m]),
)
R.output(gv)
return gv
@@ -1522,8 +1518,8 @@ def test_symbolic_var_in_buffer_shape():
def fused(
X_handle: T.handle,
Y: T.Buffer((T.int64(2048), T.int64(128)), "float32"),
- rotary_handle: T.handle,
m: T.int64,
+ rotary_handle: T.handle,
):
T.func_attr({"tirx.noalias": True})
@@ -1562,9 +1558,8 @@ def test_symbolic_var_in_buffer_shape():
with R.dataflow():
gv = R.call_tir(
cls.fused,
- (x, y),
+ (x, y, m),
out_ty=R.Tensor([1, sequence_length, 32, 128], "float32"),
- tir_vars=R.shape([m]),
)
R.output(gv)
return gv
@@ -1765,10 +1760,9 @@ def
test_symbolic_var_called_with_multiple_static_shapes():
def test_symbolic_var_called_with_static_argument():
"""A dynamic PrimFunc may accept a static argument
- The `tir_vars` parameter in `R.call_tir` contains definitions for
- all TIR variables explicitly listed in the function signature, and
- contains the TIR expression to be passed as the argument for for
- each parameter.
+ Primitive arguments in `R.call_tir` contain definitions for all TIR
+ variables explicitly listed in the function signature, and contain
+ the TIR expression to be passed for each parameter.
This test is identical to the earlier test named
"test_symbolic_var_called_with_static_shape", except for the
@@ -1780,8 +1774,8 @@ def test_symbolic_var_called_with_static_argument():
@T.prim_func(private=True, s_tir=True)
def sum_1d(
X_handle: T.handle,
- Y: T.Buffer([T.int64(1)], "float32"),
num_elements: T.int64,
+ Y: T.Buffer([T.int64(1)], "float32"),
):
X = T.match_buffer(X_handle, [num_elements], "float32")
@@ -1801,9 +1795,8 @@ def test_symbolic_var_called_with_static_argument():
with R.dataflow():
gv = R.call_tir(
cls.sum_1d,
- [x],
+ [x, 64],
out_ty=R.Tensor([1], dtype="float32"),
- tir_vars=R.shape([64]),
)
R.output(gv)
return gv
@@ -2524,7 +2517,7 @@ def test_primitive_scalar_parameter_preserves_identity():
@I.ir_module(s_tir=True)
class Before:
@T.prim_func(private=True, s_tir=True)
- def add_scalar(x: T.Buffer((4,), "int64"), y: T.Buffer((1,), "int64"),
p: T.int64):
+ def add_scalar(x: T.Buffer((4,), "int64"), p: T.int64, y:
T.Buffer((1,), "int64")):
for i in range(1):
with T.sblock("add"):
vi = T.axis.spatial(1, i)
@@ -2536,9 +2529,8 @@ def test_primitive_scalar_parameter_preserves_identity():
cls = Before
out = R.call_tir(
cls.add_scalar,
- (x,),
+ (x, p),
out_ty=R.Tensor((1,), "int64"),
- tir_vars=R.shape([p]),
)
return out
@@ -2553,5 +2545,38 @@ def test_primitive_scalar_parameter_preserves_identity():
assert tvm.tirx.analysis.verify_well_formed(after["fused"])
+def test_inplace_argument_after_primitive_scalar():
+ @I.ir_module(s_tir=True)
+ class Before:
+ @T.prim_func(private=True, s_tir=True)
+ def add_scalar_inplace(p: T.int64, x: T.Buffer((4,), "int64")):
+ for i in range(4):
+ with T.sblock("add"):
+ vi = T.axis.spatial(4, i)
+ x[vi] = x[vi] + p
+
+ @R.function(private=True)
+ def fused(p: R.Prim("int64"), x: R.Tensor((4,), "int64")) ->
R.Tensor((4,), "int64"):
+ R.func_attr({"Primitive": True})
+ cls = Before
+ out = R.call_tir_inplace(
+ cls.add_scalar_inplace,
+ (p, x),
+ inplace_indices=[1],
+ out_ty=R.Tensor((4,), "int64"),
+ )
+ return out
+
+ @R.function
+ def main(p: R.Prim("int64"), x: R.Tensor((4,), "int64")) ->
R.Tensor((4,), "int64"):
+ cls = Before
+ out = cls.fused(p, x)
+ return out
+
+ after = relax.transform.FuseTIR()(Before)
+ assert relax.analysis.check_well_formed(after)
+ assert tvm.tirx.analysis.verify_well_formed(after["fused"])
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_lazy_transform_params.py
b/tests/python/relax/test_transform_lazy_transform_params.py
index 0ff259b948..5c1e602cad 100644
--- a/tests/python/relax/test_transform_lazy_transform_params.py
+++ b/tests/python/relax/test_transform_lazy_transform_params.py
@@ -422,8 +422,7 @@ def test_lazy_transform_params_with_symbolic_vars():
param = params[0]
transformed = R.call_tir(
cls.slice_buffer,
- (param,),
- tir_vars=[slice_index],
+ (param, slice_index),
out_ty=R.Tensor((16,), dtype="float32"),
)
output = (transformed,)
@@ -432,8 +431,8 @@ def test_lazy_transform_params_with_symbolic_vars():
@T.prim_func(private=True, s_tir=True)
def slice_buffer(
Input: T.Buffer((16, 16), "float32"),
- Output: T.Buffer(16, "float32"),
slice_index: T.int64,
+ Output: T.Buffer(16, "float32"),
):
for i in T.grid(16):
with T.sblock("slice_buffer"):
@@ -455,8 +454,7 @@ def test_lazy_transform_params_with_symbolic_vars():
param_m: R.Tensor((16, 16), dtype="float32") = gv
transformed = R.call_tir(
cls.slice_buffer,
- (param_m,),
- tir_vars=[slice_index],
+ (param_m, slice_index),
out_ty=R.Tensor((16,), dtype="float32"),
)
unused_1_ = R.vm.kill_object(param_m)
@@ -468,8 +466,8 @@ def test_lazy_transform_params_with_symbolic_vars():
@T.prim_func(private=True, s_tir=True)
def slice_buffer(
Input: T.Buffer((16, 16), "float32"),
- Output: T.Buffer(16, "float32"),
slice_index: T.int64,
+ Output: T.Buffer(16, "float32"),
):
for i in T.grid(16):
with T.sblock("slice_buffer"):
diff --git a/tests/python/relax/test_transform_legalize_ops_create_datatype.py
b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
index 255da3671d..c76726fffd 100644
--- a/tests/python/relax/test_transform_legalize_ops_create_datatype.py
+++ b/tests/python/relax/test_transform_legalize_ops_create_datatype.py
@@ -600,11 +600,11 @@ def test_arange_symbolic():
def main(x: R.Tensor(["n"], "float32")):
cls = Expected
n = T.int64()
- gv = R.call_tir(cls.arange, R.tuple(), out_ty=R.Tensor((n // 2,),
dtype="int64"), tir_vars=R.shape([n]))
+ gv = R.call_tir(cls.arange, (n,), out_ty=R.Tensor((n // 2,),
dtype="int64"))
return gv
@T.prim_func(private=True, s_tir=True)
- def arange(var_T_arange: T.handle, n: T.int64):
+ def arange(n: T.int64, var_T_arange: T.handle):
T.func_attr({"tirx.noalias": True})
T_arange = T.match_buffer(var_T_arange, (n // T.int64(2),),
"int64")
for ax0 in range(n // T.int64(2)):
@@ -665,11 +665,11 @@ def test_shape_to_tensor_symbolic():
n = T.int64()
cls = Expected
gv: R.Shape([m, n]) = R.shape_of(x)
- gv_1 = R.call_tir(cls.shape_to_tensor, R.tuple(),
out_ty=R.Tensor((2,), dtype="int64"), tir_vars=R.shape([m, n]))
+ gv_1 = R.call_tir(cls.shape_to_tensor, (m, n),
out_ty=R.Tensor((2,), dtype="int64"))
return gv_1
@T.prim_func(private=True, s_tir=True)
- def shape_to_tensor(shape_to_tensor: T.Buffer((T.int64(2),), "int64"),
m: T.int64, n: T.int64):
+ def shape_to_tensor(m: T.int64, n: T.int64, shape_to_tensor:
T.Buffer((T.int64(2),), "int64")):
T.func_attr({"tirx.noalias": True})
for i in range(T.int64(2)):
with T.sblock("shape_to_tensor"):
@@ -697,11 +697,11 @@ def test_shape_to_tensor_mixed():
m = T.int64()
cls = Expected
gv: R.Shape([m, 3]) = R.shape_of(x)
- gv_1 = R.call_tir(cls.shape_to_tensor, R.tuple(),
out_ty=R.Tensor((2,), dtype="int64"), tir_vars=R.shape([m]))
+ gv_1 = R.call_tir(cls.shape_to_tensor, (m,), out_ty=R.Tensor((2,),
dtype="int64"))
return gv_1
@T.prim_func(private=True, s_tir=True)
- def shape_to_tensor(shape_to_tensor: T.Buffer((T.int64(2),), "int64"),
m: T.int64):
+ def shape_to_tensor(m: T.int64, shape_to_tensor:
T.Buffer((T.int64(2),), "int64")):
T.func_attr({"tirx.noalias": True})
for i in range(T.int64(2)):
with T.sblock("shape_to_tensor"):
diff --git a/tests/python/relax/test_transform_legalize_ops_distributed.py
b/tests/python/relax/test_transform_legalize_ops_distributed.py
index 83338570c3..2c7deb039b 100644
--- a/tests/python/relax/test_transform_legalize_ops_distributed.py
+++ b/tests/python/relax/test_transform_legalize_ops_distributed.py
@@ -37,7 +37,7 @@ def test_redistribute_replica_to_shard():
@I.ir_module(s_tir=True)
class Expected:
@T.prim_func(private=True, s_tir=True)
- def strided_slice(A: T.Buffer((T.int64(10), T.int64(10)), "float32"),
redistribute_replica_to_shard: T.Buffer((T.int64(10), T.int64(5)), "float32"),
worker_id: T.int64):
+ def strided_slice(A: T.Buffer((T.int64(10), T.int64(10)), "float32"),
worker_id: T.int64, redistribute_replica_to_shard: T.Buffer((T.int64(10),
T.int64(5)), "float32")):
T.func_attr({"tirx.noalias": True})
# with T.sblock("root"):
for i0, i1 in T.grid(T.int64(10), T.int64(5)):
@@ -53,7 +53,7 @@ def test_redistribute_replica_to_shard():
cls = Expected
gv: R.Shape(ndim=-1) =
R.call_pure_packed("runtime.disco.worker_id", ty_args=(R.Shape(ndim=-1),))
gv1: R.Shape([worker_id]) = R.match_cast(gv, R.Shape([worker_id]))
- gv0 = R.call_tir(cls.strided_slice, (x,), out_ty=R.Tensor((10, 5),
dtype="float32"), tir_vars=R.shape([worker_id]))
+ gv0 = R.call_tir(cls.strided_slice, (x, worker_id),
out_ty=R.Tensor((10, 5), dtype="float32"))
return gv0
# fmt: on
diff --git
a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
index 549e4e5d99..a026507ae9 100644
--- a/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
+++ b/tests/python/relax/test_transform_legalize_ops_index_linear_algebra.py
@@ -71,11 +71,11 @@ def test_take_prim_value():
class Expected:
@R.function
def main(x: R.Tensor((2, 3, 4), "float32"), index: R.Prim("int64")) ->
R.Tensor((2, 4), "float32"):
- gv = R.call_tir(Expected.take, (x,), R.Tensor((2, 4),
dtype="float32"), tir_vars=R.shape([index]))
+ gv = R.call_tir(Expected.take, (x, index), R.Tensor((2, 4),
dtype="float32"))
return gv
@T.prim_func(private=True, s_tir=True)
- def take(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)),
"float32"), T_take: T.Buffer((T.int64(2), T.int64(4)), "float32"), index:
T.int64):
+ def take(rxplaceholder: T.Buffer((T.int64(2), T.int64(3), T.int64(4)),
"float32"), index: T.int64, T_take: T.Buffer((T.int64(2), T.int64(4)),
"float32")):
T.func_attr({"tirx.noalias": True})
for i0, i2 in T.grid(T.int64(2), T.int64(4)):
with T.sblock("T_take"):
diff --git a/tests/python/relax/test_transform_legalize_ops_manipulate.py
b/tests/python/relax/test_transform_legalize_ops_manipulate.py
index 7a09fd75da..250e953f64 100644
--- a/tests/python/relax/test_transform_legalize_ops_manipulate.py
+++ b/tests/python/relax/test_transform_legalize_ops_manipulate.py
@@ -881,11 +881,11 @@ def test_split_by_indices_n_section_divisible_symbolic():
def main(dumb_param: R.Tensor(("n",)), x: R.Tensor(("m", "(n * 3)"),
"float32")) -> R.Tuple(R.Tensor(("m", "((n * 3) // 3)"), "float32"),
R.Tensor(("m", "((((n * 3) // 3) * 2) - ((n * 3) // 3))"), "float32"),
R.Tensor(("m", "((n * 3) - (((n * 3) // 3) * 2))"), "float32")):
m = T.int64()
n = T.int64()
- gv = R.call_tir(Expected.split, (x,), [R.Tensor((m, ((n * 3 + 3 -
1) // 3)), "float32"), R.Tensor((m, ((((n * 3 + 3 - 1) // 3) * 2) - ((n * 3 + 3
- 1) // 3))), "float32"), R.Tensor((m, ((n * 3) - (((n * 3 + 3 - 1) // 3) *
2))), "float32")], tir_vars=R.shape([n]))
+ gv = R.call_tir(Expected.split, (x, n), [R.Tensor((m, ((n * 3 + 3
- 1) // 3)), "float32"), R.Tensor((m, ((((n * 3 + 3 - 1) // 3) * 2) - ((n * 3 +
3 - 1) // 3))), "float32"), R.Tensor((m, ((n * 3) - (((n * 3 + 3 - 1) // 3) *
2))), "float32")])
return gv
@T.prim_func(private=True, s_tir=True)
- def split(var_rxplaceholder: T.handle, var_T_split_sections: T.handle,
var_T_split_sections_1: T.handle, var_T_split_sections_2: T.handle, n: T.int64):
+ def split(var_rxplaceholder: T.handle, n: T.int64,
var_T_split_sections: T.handle, var_T_split_sections_1: T.handle,
var_T_split_sections_2: T.handle):
T.func_attr({"tirx.noalias": True})
m = T.int64()
rxplaceholder = T.match_buffer(var_rxplaceholder, [m, n *
T.int64(3)], dtype="float32")
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py
b/tests/python/relax/test_transform_legalize_ops_nn.py
index 058a4e6081..c70ea1b670 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -3186,7 +3186,7 @@ def test_group_norm_symbolic():
@tvm.script.ir_module
class Expected:
@T.prim_func(private=True, s_tir=True)
- def group_norm(var_rxplaceholder: T.handle, var_rxplaceholder_1:
T.handle, var_rxplaceholder_2: T.handle, var_T_reshape: T.handle, c: T.int64):
+ def group_norm(var_rxplaceholder: T.handle, var_rxplaceholder_1:
T.handle, var_rxplaceholder_2: T.handle, c: T.int64, var_T_reshape: T.handle):
T.func_attr({"tirx.noalias": True})
n = T.int64()
h = T.int64()
@@ -3251,7 +3251,7 @@ def test_group_norm_symbolic():
c = T.int64()
h = T.int64()
w = T.int64()
- gv = R.call_tir(Expected.group_norm, (x, gamma, beta),
out_ty=R.Tensor((n, 4 * c, h, w), dtype="float32"), tir_vars=R.shape([c]))
+ gv = R.call_tir(Expected.group_norm, (x, gamma, beta, c),
out_ty=R.Tensor((n, 4 * c, h, w), dtype="float32"))
return gv
# fmt: on
mod = LegalizeOps()(GroupNorm)
diff --git a/tests/python/relax/test_transform_lift_transform_params.py
b/tests/python/relax/test_transform_lift_transform_params.py
index 40fefd26a8..37b3516b60 100644
--- a/tests/python/relax/test_transform_lift_transform_params.py
+++ b/tests/python/relax/test_transform_lift_transform_params.py
@@ -1505,14 +1505,12 @@ def test_symbolic_var_from_shape():
with R.dataflow():
B_slice = R.call_tir(
cls.slice,
- [B],
- tir_vars=R.ShapeExpr([slice_index]),
+ [B, slice_index],
out_ty=R.Tensor([16], dtype="int32"),
)
A_slice = R.call_tir(
cls.slice,
- [A],
- tir_vars=R.ShapeExpr([slice_index]),
+ [A, slice_index],
out_ty=R.Tensor([16], dtype="int32"),
)
A_scale = R.multiply(A_slice, B_slice)
@@ -1522,8 +1520,8 @@ def test_symbolic_var_from_shape():
@T.prim_func(private=True, s_tir=True)
def slice(
Input_2d: T.Buffer(shape=[16, 16], dtype="int32"),
- Output_Slice: T.Buffer(shape=[16], dtype="int32"),
slice_index: T.int64,
+ Output_Slice: T.Buffer(shape=[16], dtype="int32"),
):
T.func_attr({"tirx.noalias": True})
for j in range(16):
@@ -1545,8 +1543,7 @@ def test_symbolic_var_from_shape():
with R.dataflow():
A_slice = R.call_tir(
cls.slice,
- [A],
- tir_vars=R.ShapeExpr([slice_index]),
+ [A, slice_index],
out_ty=R.Tensor([16], dtype="int32"),
)
A_scale = R.multiply(A_slice, B_slice)
@@ -1565,8 +1562,7 @@ def test_symbolic_var_from_shape():
# extra_symbolic_vars = params[1]
B_slice = R.call_tir(
cls.slice,
- [B],
- tir_vars=R.ShapeExpr([slice_index]),
+ [B, slice_index],
out_ty=R.Tensor([16], dtype="int32"),
)
output = (R.ShapeExpr([slice_index]), B_slice)
@@ -1576,8 +1572,8 @@ def test_symbolic_var_from_shape():
@T.prim_func(private=True, s_tir=True)
def slice(
Input_2d: T.Buffer(shape=[16, 16], dtype="int32"),
- Output_Slice: T.Buffer(shape=[16], dtype="int32"),
slice_index: T.int64,
+ Output_Slice: T.Buffer(shape=[16], dtype="int32"),
):
T.func_attr({"tirx.noalias": True})
for j in range(16):
diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
index ef9ef115b1..bbe22389ac 100644
--- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
+++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py
@@ -659,8 +659,7 @@ def test_rewrite_static_reshape():
# y = R.reshape(x, R.shape([N // 4, 4]))
# z = R.call_tir(
# cls.add,
-# (y, y),
-# tir_vars=[N],
+# (y, y, N),
# out_ty=R.Tensor((N // 4, 4), dtype="float32"),
# )
# R.output(z)
@@ -670,8 +669,8 @@ def test_rewrite_static_reshape():
# def add(
# y1_handle: T.handle,
# y2_handle: T.handle,
-# z_handle: T.handle,
# N: T.int64,
+# z_handle: T.handle,
# ):
# y1 = T.match_buffer(y1_handle, [N // 4, 4], "float32")
@@ -724,8 +723,7 @@ def test_rewrite_dynamic_reshape():
y = R.reshape(x, R.shape([N * 4, T.int64(4)]))
z = R.call_tir(
cls.add,
- (y, y),
- tir_vars=[N],
+ (y, y, N),
out_ty=R.Tensor((N * 4, 4), dtype="float32"),
)
R.output(z)
@@ -735,8 +733,8 @@ def test_rewrite_dynamic_reshape():
def add(
y1_handle: T.handle,
y2_handle: T.handle,
- z_handle: T.handle,
N: T.int64,
+ z_handle: T.handle,
):
y1 = T.match_buffer(y1_handle, [N * 4, T.int64(4)], "float32")
y2 = T.match_buffer(y2_handle, [N * 4, T.int64(4)], "float32")
diff --git a/tests/python/relax/test_tvmscript_parser.py
b/tests/python/relax/test_tvmscript_parser.py
index 50628970e7..b91a1e4544 100644
--- a/tests/python/relax/test_tvmscript_parser.py
+++ b/tests/python/relax/test_tvmscript_parser.py
@@ -982,11 +982,11 @@ def test_call_tir_with_tir_var():
) -> R.Tensor(("n * 2",), "float32"):
n = T.int64()
cls = Module
- y = R.call_tir(cls.copy, x, R.Tensor((n * 2,), dtype="float32"),
tir_vars=(n,))
+ y = R.call_tir(cls.copy, (x, n), R.Tensor((n * 2,),
dtype="float32"))
return y
@T.prim_func(s_tir=True)
- def copy(var_x: T.handle, var_y: T.handle, n: T.int64):
+ def copy(var_x: T.handle, n: T.int64, var_y: T.handle):
X = T.match_buffer(var_x, (n * 2,), dtype="float32")
Y = T.match_buffer(var_y, (n * 2,), dtype="float32")
for i in T.grid(n * 2):
diff --git a/tests/python/relax/test_tvmscript_printer_relax.py
b/tests/python/relax/test_tvmscript_printer_relax.py
index c47d9847b0..f2aeaf9b1a 100644
--- a/tests/python/relax/test_tvmscript_printer_relax.py
+++ b/tests/python/relax/test_tvmscript_printer_relax.py
@@ -485,14 +485,14 @@ def test_shape_expr():
def test_call():
x = tirx.Var("x", "int64")
a = relax.Var("a", relax.TensorType([1, x, 3], "float32"))
- o0 = relax.call_tir(relax.GlobalVar("tir_func"), args=a, out_ty=a.ty,
tir_vars=[x])
+ o0 = relax.call_tir(relax.GlobalVar("tir_func"), args=(a, x), out_ty=a.ty)
o1 = relax.call_dps_packed("my_dps_func", args=a, out_ty=a.ty)
_assert_print(
o0,
"""
x = T.int64()
a: R.Tensor((1, x, 3), dtype="float32")
-R.call_tir(tir_func, (a,), out_ty=R.Tensor((1, x, 3), dtype="float32"),
tir_vars=R.shape([x]))
+R.call_tir(tir_func, (a, x), out_ty=R.Tensor((1, x, 3), dtype="float32"))
""",
)
_assert_print(
@@ -535,10 +535,10 @@ def test_call_tir_inplace():
(
x,
y,
+ t,
),
inplace_indices=[-1, 0],
out_ty=[R.Tensor((32, 32), dtype="int32"), R.Tensor((32, 32),
dtype="int32")],
- tir_vars=[t],
)
_assert_print(
call,
@@ -546,7 +546,7 @@ def test_call_tir_inplace():
x: R.Tensor((32, 32), dtype="int32")
y: R.Tensor((32, 32), dtype="int32")
t = T.int64()
-R.call_tir_inplace(tir_func, (x, y), out_ty=[R.Tensor((32, 32),
dtype="int32"), R.Tensor((32, 32), dtype="int32")], inplace_indices=[-1, 0],
tir_vars=R.shape([t]))
+R.call_tir_inplace(tir_func, (x, y, t), out_ty=[R.Tensor((32, 32),
dtype="int32"), R.Tensor((32, 32), dtype="int32")], inplace_indices=[-1, 0])
""",
)
