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 14dcefa2b4 [REFACTOR][TIRx] Store tile operands and configuration as
expressions (#20576)
14dcefa2b4 is described below
commit 14dcefa2b446d1a9c6038846c0f379ca7c7a2477
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 7 07:20:12 2026 -0400
[REFACTOR][TIRx] Store tile operands and configuration as expressions
(#20576)
---
include/tvm/ir/base_expr.h | 1 +
include/tvm/tirx/tile_primitive.h | 21 +-
.../backend/cuda/tile_primitive/copy/vec_forced.py | 6 +-
.../cuda/tile_primitive/copy_async/tcgen05_cp.py | 11 +-
.../backend/cuda/tile_primitive/copy_async/tma.py | 18 +-
.../cuda/tile_primitive/elementwise/ops/binary.py | 2 +-
.../cuda/tile_primitive/elementwise/ops/unary.py | 30 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 21 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 4 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 28 +-
.../trn/tile_primitive/instruction_generator.py | 6 +-
.../backend/trn/tile_primitive/private_alloc.py | 25 +-
.../backend/trn/tile_primitive/reduction/utils.py | 2 +-
.../trn/tile_primitive/unary/with_bias_scale.py | 12 +-
python/tvm/tirx/operator/tile_primitive/ops.py | 83 +++--
python/tvm/tirx/script/ir_builder/tirx.py | 381 +++++++++++----------
python/tvm/tirx/script/tile.py | 15 +-
python/tvm/tirx/tile_primitive.py | 17 +-
src/backend/trn/transform/lower_trainium_layout.cc | 3 +-
src/ir/expr.cc | 13 +
src/relax/op/tensor/create.cc | 3 +-
src/relax/op/tensor/create.h | 2 +-
src/relax/op/tensor/manipulate.cc | 5 +-
src/relax/op/tensor/manipulate.h | 2 +-
src/tirx/ir/stmt_functor.cc | 85 +----
src/tirx/ir/tir_visitor_with_path.cc | 29 +-
src/tirx/ir/tirx_stmt.cc | 28 +-
src/tirx/op/tirx.cc | 60 +++-
src/tirx/script/printer/stmt.cc | 44 ++-
src/tirx/transform/lower_tirx_cleanup.cc | 4 +-
.../tile_primitive/cuda/elementwise/test_unary.py | 56 +--
.../tile_primitive/trn/test_compose_op_trn.py | 4 +-
.../tile_primitive/trn/test_private_alloc_trn.py | 14 +-
.../operator/tile_primitive/trn/test_unary_trn.py | 6 +-
tests/python/tirx/script/test_tirx_script_gpu.py | 2 +-
tests/python/tirx/test_hint.py | 2 +-
36 files changed, 606 insertions(+), 439 deletions(-)
diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h
index 74e9c7a65f..b8760f673b 100644
--- a/include/tvm/ir/base_expr.h
+++ b/include/tvm/ir/base_expr.h
@@ -621,6 +621,7 @@ template <>
struct TypeTraits<Expr> : public ObjectRefWithFallbackTraitsBase<Expr,
PrimExpr, ffi::String> {
TVM_FFI_INLINE static Expr ConvertFallbackValue(PrimExpr value) { return
value; }
TVM_DLL static Expr ConvertFallbackValue(ffi::String value);
+ TVM_DLL static std::optional<Expr> TryCastFromAnyView(const TVMFFIAny* src);
};
} // namespace ffi
diff --git a/include/tvm/tirx/tile_primitive.h
b/include/tvm/tirx/tile_primitive.h
index 1716e29a03..7af00985c1 100644
--- a/include/tvm/tirx/tile_primitive.h
+++ b/include/tvm/tirx/tile_primitive.h
@@ -148,9 +148,9 @@ class TilePrimitiveCallNode : public StmtNode {
public:
explicit TilePrimitiveCallNode(ffi::UnsafeInit tag) : op(tag) {}
- TilePrimitiveCallNode(tvm::Op op, ffi::Array<ffi::Any> args,
+ TilePrimitiveCallNode(tvm::Op op, ffi::Array<Expr> args,
ffi::Map<ffi::String, TensorVar> workspace,
- ffi::Map<ffi::String, ffi::Any> config,
ffi::Optional<ffi::String> dispatch,
+ ffi::Map<ffi::String, Expr> config,
ffi::Optional<ffi::String> dispatch,
ExecScope scope)
: op(std::move(op)),
args(std::move(args)),
@@ -163,13 +163,13 @@ class TilePrimitiveCallNode : public StmtNode {
tvm::Op op;
// Arguments to the operator.
- ffi::Array<ffi::Any> args;
+ ffi::Array<Expr> args;
// Workspace (pre-allocated buffers) for the operator.
ffi::Map<ffi::String, TensorVar> workspace;
// Config for the operator/scheduler.
- ffi::Map<ffi::String, ffi::Any> config;
+ ffi::Map<ffi::String, Expr> config;
// Optional dispatch variant name registered via @register_dispatch.
ffi::Optional<ffi::String> dispatch{std::nullopt};
@@ -197,14 +197,12 @@ class TilePrimitiveCallNode : public StmtNode {
*/
class TilePrimitiveCall : public Stmt {
public:
- TVM_DLL TilePrimitiveCall(tvm::Op op, ffi::Array<ffi::Any> args,
+ TVM_DLL TilePrimitiveCall(tvm::Op op, ffi::Array<Expr> args,
ffi::Map<ffi::String, TensorVar> workspace = {},
- ffi::Map<ffi::String, ffi::Any> config = {},
+ ffi::Map<ffi::String, Expr> config = {},
ffi::Optional<ffi::String> dispatch = std::nullopt,
ExecScope scope = ExecScope(ScopeKind::kThread));
- static bool IsValidOpCallArgType(const ffi::Any& arg);
-
explicit TilePrimitiveCall(ffi::ObjectPtr<TilePrimitiveCallNode> node) :
Stmt(std::move(node)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TilePrimitiveCall, Stmt,
TilePrimitiveCallNode);
@@ -257,10 +255,16 @@ TVM_DLL const Op& gemm_async();
TVM_DLL const Op& zero();
TVM_DLL const Op& sqrt();
+TVM_DLL const Op& sqrt_with_scale_bias();
TVM_DLL const Op& exp();
+TVM_DLL const Op& exp_with_scale_bias();
TVM_DLL const Op& exp2();
+TVM_DLL const Op& exp2_with_scale_bias();
+
+TVM_DLL const Op& log2();
+TVM_DLL const Op& log2_with_scale_bias();
TVM_DLL const Op& add();
@@ -289,6 +293,7 @@ TVM_DLL const Op& reduce_negate();
TVM_DLL const Op& binary_reduce();
TVM_DLL const Op& unary_reduce();
+TVM_DLL const Op& unary_reduce_with_scale_bias();
TVM_DLL const Op& binary_chain();
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
index b9e392bae5..a5cc788b91 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
@@ -80,12 +80,12 @@ def _ld_cache_config(op_call: TilePrimitiveCall) ->
tuple[str | None, dict[str,
"""
cache = op_call.config.get("cache", None)
if cache is not None:
- cache = str(cache)
+ cache = cache.value
hints: dict[str, str] = {}
for key in _LD_CACHE_HINT_KEYS:
value = op_call.config.get(key, None)
- if value is not None and str(value):
- hints[key] = str(value)
+ if value is not None and value.value:
+ hints[key] = value.value
return cache, hints
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
index 0a62e645ea..aed613780e 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
@@ -196,8 +196,10 @@ def _cp_lane_replica_pattern(shape: str, multicast: str):
def _resolve_cp_shape(op_call: TilePrimitiveCall):
"""Resolve (shape, multicast) from an explicit ``shape=`` config."""
- shape = str(op_call.config["shape"])
+ shape = op_call.config["shape"].value
multicast = op_call.config.get("multicast")
+ if isinstance(multicast, tvm.ir.StringImm):
+ multicast = multicast.value
allowed = _CP_SHAPE_MULTICASTS.get(shape)
if allowed is None:
raise ValueError(
@@ -334,6 +336,8 @@ def _build_plan(op_call: TilePrimitiveCall):
return _plan_for_shape(op_call, shape, multicast)
# No shape config: infer from the buffer layouts.
multicast_cfg = op_call.config.get("multicast")
+ if isinstance(multicast_cfg, tvm.ir.StringImm):
+ multicast_cfg = multicast_cfg.value
errors = []
for shape, multicast in _CP_SHAPE_CANDIDATES:
if multicast_cfg is not None and str(multicast_cfg) != multicast:
@@ -740,7 +744,10 @@ def _validate_smem_tmem_copy(op_call: TilePrimitiveCall,
sctx: DispatchContext):
# need synchronization.
# -----------------------------------------------------------------------------
def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) ->
PrimFunc | None:
- if op_call.config.get("decompress"):
+ decompress = op_call.config.get("decompress")
+ if isinstance(decompress, StringImm):
+ decompress = decompress.value
+ if decompress:
# fp4/fp6->fp8 in-flight decompression needs dtype-pair plan derivation
# the planner can't lower; reject loudly rather than copy
undecompressed.
raise ValueError("tcgen05.cp planner does not support decompress")
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
index 1913a423b7..dab89ef5d7 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
@@ -339,6 +339,8 @@ def _target_sm(arch: str) -> int:
def _normalize_l2_promotion(value) -> int:
+ if isinstance(value, tvm.ir.StringImm):
+ value = value.value
if value is None:
return 2
if isinstance(value, IntImm):
@@ -360,6 +362,8 @@ def _normalize_l2_promotion(value) -> int:
def _normalize_oob(value) -> int:
+ if isinstance(value, tvm.ir.StringImm):
+ value = value.value
if value is None or value == "zero":
return 0
if value == "nan":
@@ -368,6 +372,8 @@ def _normalize_oob(value) -> int:
def _normalize_cache_hint(cache_hint):
+ if isinstance(cache_hint, tvm.ir.StringImm):
+ cache_hint = cache_hint.value
if cache_hint is None:
return "", None
if isinstance(cache_hint, str):
@@ -1585,6 +1591,8 @@ def _runtime_config(op_call, sctx, direction: str, *,
explicit: bool):
fail("cta_mask is only valid for global-to-shared TMA")
use_tma_reduce = op_call.config.get("use_tma_reduce")
+ if isinstance(use_tma_reduce, tvm.ir.StringImm):
+ use_tma_reduce = use_tma_reduce.value
if use_tma_reduce is not None:
if direction != "s2g":
fail("use_tma_reduce is only valid for shared-to-global TMA")
@@ -1612,7 +1620,7 @@ def _runtime_config(op_call, sctx, direction: str, *,
explicit: bool):
"l2_promotion":
_normalize_l2_promotion(op_call.config.get("tensormap_l2_promotion")),
"oob_fill": _normalize_oob(op_call.config.get("oob")),
"prefetch": bool(op_call.config.get("prefetch_tensormap", False)),
- "tma_dtype": op_call.config.get("tma_dtype"),
+ "tma_dtype": (op_call.config["tma_dtype"].value if "tma_dtype" in
op_call.config else None),
"target_arch": sctx.target.arch,
}
@@ -1860,7 +1868,7 @@ def _explicit_spec_for_gmem(
def _normalize_gather4(value):
if value is None:
return ()
- if not isinstance(value, list | tuple | tvm.ir.Array) or len(value) != 4:
+ if not isinstance(value, list | tuple | tvm.ir.Array | tvm.ir.Tuple) or
len(value) != 4:
fail("tma_explicit gather4 must contain exactly four row coordinates")
return tuple(value)
@@ -1899,11 +1907,11 @@ def _validate_gather4_dst(s_buf: Var, s_starts,
s_extents, spec: TensorMapSpec)
def _normalize_src_selector(value):
if value is None:
return ()
- if not isinstance(value, list | tuple | tvm.ir.Array):
- fail("tma_explicit src_selector must be a list of (condition, global
Var/view)")
+ if not isinstance(value, list | tuple | tvm.ir.Array | tvm.ir.Tuple) or
len(value) == 0:
+ fail("tma_explicit src_selector must be a nonempty list of (condition,
global Var/view)")
result = []
for idx, item in enumerate(value):
- if not isinstance(item, list | tuple | tvm.ir.Array) or len(item) != 2:
+ if not isinstance(item, list | tuple | tvm.ir.Array | tvm.ir.Tuple) or
len(item) != 2:
fail(f"tma_explicit src_selector[{idx}] must be a (condition,
Var/view) pair")
condition, buffer = item
if not isinstance(condition, tvm.tirx.Expr):
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
index 3b943440c5..1d3a214842 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
@@ -82,7 +82,7 @@ def _parse_binary_for(op_name: str):
extras: dict[str, Any] = {}
rm = op.config.get("rounding_mode", None)
if rm is not None:
- extras["rounding_mode"] = rm
+ extras["rounding_mode"] = rm.value
return Plan(dst=_dst, srcs=srcs, extras=extras), None
return parse
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
index ed6d921414..ee1458b5e8 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
@@ -17,8 +17,8 @@
"""Unary elementwise ops: zero / fill / reciprocal / sqrt / exp / exp2 / log2
/ silu.
-All carry the same ``T.<unary>(dst, src[, bias, scale])`` shape (bias / scale
-optional; ``silu`` ignores bias/scale to preserve legacy behavior).
+Plain unary operations take dst and src. Explicit scale/bias variants
additionally
+take scale and bias, applying both before the unary operation.
"""
from __future__ import annotations
@@ -28,18 +28,22 @@ from typing import Any
from tvm.ir import TensorRegion, is_prim_expr
from tvm.script import tirx as T
from tvm.tirx import TilePrimitiveCall
-from tvm.tirx.expr import FloatImm
from .._common import scalar_dtype
from . import OpSpec, Plan, SrcSpec
def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]:
- """T.<unary>(dst, src[, bias, scale]) → Plan."""
+ """T.<unary>(dst, src[, scale, bias]) → Plan."""
_dst: TensorRegion = op.args[0]
_src = op.args[1]
- _bias = op.args[2] if len(op.args) > 2 else None
- _scale = op.args[3] if len(op.args) > 2 else None
+ _scale = op.args[2] if len(op.args) == 4 else None
+ _bias = op.args[3] if len(op.args) == 4 else None
+
+ if _bias is not None and not isinstance(_bias, TensorRegion) and not
is_prim_expr(_bias):
+ return None, f"unsupported bias type {type(_bias).__name__}"
+ if _scale is not None and not is_prim_expr(_scale):
+ return None, f"unsupported scale type {type(_scale).__name__}"
srcs: list[SrcSpec] = []
if isinstance(_src, TensorRegion):
@@ -51,7 +55,7 @@ def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None,
str | None]:
extras: dict[str, Any] = {
"scale": _scale,
- "bias_const": _bias if isinstance(_bias, FloatImm) else None,
+ "bias_const": _bias if is_prim_expr(_bias) else None,
}
if isinstance(_bias, TensorRegion):
srcs.append(SrcSpec(buf_region=_bias))
@@ -115,8 +119,20 @@ UNARY_OPS: dict[str, OpSpec] = {
"fill": OpSpec("fill", _parse_unary, _compute_fill, _check_unary_extras),
"reciprocal": OpSpec("reciprocal", _parse_unary, _compute_reciprocal,
_check_unary_extras),
"sqrt": OpSpec("sqrt", _parse_unary, _with_bias_scale(T.sqrt),
_check_unary_extras),
+ "sqrt_with_scale_bias": OpSpec(
+ "sqrt_with_scale_bias", _parse_unary, _with_bias_scale(T.sqrt),
_check_unary_extras
+ ),
"exp": OpSpec("exp", _parse_unary, _with_bias_scale(T.exp),
_check_unary_extras),
+ "exp_with_scale_bias": OpSpec(
+ "exp_with_scale_bias", _parse_unary, _with_bias_scale(T.exp),
_check_unary_extras
+ ),
"exp2": OpSpec("exp2", _parse_unary, _with_bias_scale(T.exp2),
_check_unary_extras),
+ "exp2_with_scale_bias": OpSpec(
+ "exp2_with_scale_bias", _parse_unary, _with_bias_scale(T.exp2),
_check_unary_extras
+ ),
"log2": OpSpec("log2", _parse_unary, _with_bias_scale(T.log2),
_check_unary_extras),
+ "log2_with_scale_bias": OpSpec(
+ "log2_with_scale_bias", _parse_unary, _with_bias_scale(T.log2),
_check_unary_extras
+ ),
"silu": OpSpec("silu", _parse_unary, _compute_silu, _check_unary_extras),
}
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
index 0c9822f94a..0e7e07669f 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -30,7 +30,7 @@ from tvm.ir import Call, DataTypeImm, DictAttrs, StringImm,
TensorRegion, Tuple
from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
-from tvm.tirx import PrimFunc
+from tvm.tirx import IntImm, PrimFunc
from tvm.tirx import op as tirx_op
from tvm.tirx.layout import (
ComposeLayout,
@@ -279,12 +279,12 @@ def _get_explicit_mma_tile(config):
values = []
for name in ("mma_m", "mma_n"):
value = config[name]
- if isinstance(value, bool):
+ if isinstance(value, bool) or (isinstance(value, IntImm) and
str(value.ty.dtype) == "bool"):
raise ValueError(
f"gemm_async[tcgen05]: {name} must be a positive integer, got
{value!r}"
)
try:
- value = operator.index(value)
+ value = int(value) if isinstance(value, IntImm) else
operator.index(value)
except TypeError as err:
raise ValueError(
f"gemm_async[tcgen05]: {name} must be a positive integer, got
{value!r}"
@@ -466,6 +466,8 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
if is_block_scaled:
SFA_buffer_region, SFB_buffer_region = op_call.sfa, op_call.sfb
transA, transB, accum = op_call.transA, op_call.transB, op_call.accum
+ transA = bool(transA) if isinstance(transA, IntImm) else transA
+ transB = bool(transB) if isinstance(transB, IntImm) else transB
SFA_buffer: tvm.ir.Var = SFA_buffer_region.source
SFB_buffer: tvm.ir.Var = SFB_buffer_region.source
SFA_scope, SFB_scope = SFA_buffer.scope(), SFB_buffer.scope()
@@ -499,8 +501,12 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
_validate_sf_tmem_layout(SFB_slice_layout, SFB_rows, SFB_K_total,
sfb_sf_mma_k, "SFB")
else:
transA, transB, accum = op_call.transA, op_call.transB, op_call.accum
+ transA = bool(transA) if isinstance(transA, IntImm) else transA
+ transB = bool(transB) if isinstance(transB, IntImm) else transB
cta_group = op_call.config.get("cta_group", 1)
+ if isinstance(cta_group, IntImm):
+ cta_group = cta_group.value
assert cta_group in [1, 2], f"tcgen05 schedule expected cta_group=1 or 2,
got {cta_group}"
# descI (pre-encoded uint32 instruction descriptor): rejected on the dense
# path (dispatcher encodes it); block-scaled callers may still pass it in.
@@ -950,7 +956,12 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
# Packed Layout-E C (M, 2, N//2) is uniquely the cta_group::1 M=64 .ws
datapath
# (PTX §9.7.16.10.5), so .ws is inferred; weight_stationary=False on it is
rejected.
if packed_n2 and not is_2x2:
- if op_call.config.get("weight_stationary") is False:
+ explicit_ws = op_call.config.get("weight_stationary")
+ if (
+ isinstance(explicit_ws, IntImm)
+ and str(explicit_ws.ty.dtype) == "bool"
+ and not int(explicit_ws)
+ ):
raise ValueError(
"gemm_async[tcgen05]: C uses the packed (M, 2, N//2):(1@TLane,
"
"64@TLane, 1@TCol) Layout-E TMEM layout, which is the M=64 "
@@ -1208,6 +1219,8 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
# smem_desc modes: "hoist" (default, encode once after alloc),
"local_hoist"
# (encode at call site, reuse via add_16B_offset), "encode"/"recompute"
(per MMA).
smem_desc_mode = op_call.config.get("smem_desc", "hoist")
+ if isinstance(smem_desc_mode, tvm.ir.StringImm):
+ smem_desc_mode = smem_desc_mode.value
local_hoist = smem_desc_mode == "local_hoist"
encode_per_mma = smem_desc_mode == "encode"
use_add = smem_desc_mode not in ("recompute", "encode")
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
index 39fb66e417..c89ce1e626 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
@@ -43,7 +43,9 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc
analyzer = init_analyzer(sctx)
# Normalize negative axes
- reduce_axes = [i if i >= 0 else len(binary_output.source.ty.shape) + i for
i in reduce_axes]
+ reduce_axes = [
+ int(i) if int(i) >= 0 else len(binary_output.source.ty.shape) + int(i)
for i in reduce_axes
+ ]
# Find instruction patterns
inst_gen = InstructionGenerator(
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
index ae8c720a3d..98a31442cc 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
@@ -39,13 +39,18 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
# Extract operation components
unary_output, reduce_output = op.dsts
- unary_input, bias, scale = op.srcs
+ unary_input = op.unary_input
+ if op.op.name == "tirx.tile.unary_reduce_with_scale_bias":
+ scale, bias = op.scale, op.bias
+ else:
+ scale, bias = 1.0, 0.0
analyzer = init_analyzer(sctx)
# Normalize axes and default values
- reduce_axes = [i if i >= 0 else len(unary_output.source.ty.shape) + i for
i in op.reduce_axes]
- scale = 1.0 if scale is None else scale
- bias = 0.0 if bias is None else bias
+ reduce_axes = [
+ int(i) if int(i) >= 0 else len(unary_output.source.ty.shape) + int(i)
+ for i in op.reduce_axes
+ ]
inst_gen = InstructionGenerator([unary_output, unary_input, bias,
reduce_output], analyzer)
reduce_dim_map = get_reduction_dim_map(unary_output, reduce_output,
reduce_axes, analyzer)
@@ -156,6 +161,21 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
return impl
+@register_dispatch(
+ "unary_reduce_with_scale_bias",
+ "trn",
+ variant="default",
+ priority=10,
+ when=[
+ predicate(
+ "exec_scope",
+ lambda op, sctx: (
+ sctx.scope_kind == "thread",
+ f"unsupported exec_scope {sctx.scope_kind}",
+ ),
+ )
+ ],
+)
@register_dispatch(
"unary_reduce",
"trn",
diff --git a/python/tvm/backend/trn/tile_primitive/instruction_generator.py
b/python/tvm/backend/trn/tile_primitive/instruction_generator.py
index 6693596dae..d7a079c097 100644
--- a/python/tvm/backend/trn/tile_primitive/instruction_generator.py
+++ b/python/tvm/backend/trn/tile_primitive/instruction_generator.py
@@ -30,7 +30,7 @@ from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import Range, TensorRegion
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
-from tvm.tirx import BufferRegion, Expr, Var, is_tensor_var
+from tvm.tirx import BufferRegion, Expr, IntImm, Var, is_tensor_var
from tvm.tirx.layout import Iter
from .dim_utils import DimensionMapper, RangeInfo, normalize_and_group
@@ -82,7 +82,9 @@ class InstructionRepr:
self.selected_data_iter_ids = selected_data_iter_ids
def bound_inst_size(self, max_inst_size: int | None, analyzer: Analyzer):
- if max_inst_size is None:
+ if max_inst_size is None or (
+ isinstance(max_inst_size, int | IntImm) and int(max_inst_size) ==
-1
+ ):
return
if analyzer.can_prove(self.size <= max_inst_size):
return
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py
b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index ceafb52120..9a7bbd260d 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -21,13 +21,15 @@ from tvm.backend.trn.tile_primitive.common import
init_analyzer, nki_dim
from tvm.backend.trn.tile_primitive.dim_utils import get_ewise_dim_map
from tvm.backend.trn.tile_primitive.instruction_generator import
InstructionGenerator
from tvm.script import tirx as T
-from tvm.tirx import FloatImm, Stmt, Var
+from tvm.tirx import FloatImm, IntImm, Stmt, Var
from tvm.tirx.operator.tile_primitive.ops import (
BinaryReduce,
Copy,
+ Exp,
Gemm,
ReduceOp,
- UnaryOpWithBiasScale,
+ Sqrt,
+ UnaryOpWithScaleBias,
UnaryReduce,
)
from tvm.tirx.operator.tile_primitive.registry import f_op_dispatcher
@@ -48,19 +50,22 @@ def _scalar_dtype(scalar) -> str:
def alloc_const_bias_trn(
op: TilePrimitiveCall, buffer_dict: dict[Any, tuple[Var, Stmt | None]],
sctx: DispatchContext
) -> dict[str, Any]:
- bias = op.bias if op.bias is not None else
FloatImm(op.dsts[0].source.ty.dtype, 0.0)
+ bias = getattr(op, "bias", FloatImm(op.dsts[0].source.ty.dtype, 0.0))
if "const_bias" in op.workspace:
return {}
if not isinstance(bias, (FloatImm)):
return {}
+ bias_key = ("const_bias", _scalar_dtype(bias), float(bias.value).hex())
par_size = op.dsts[0].source.ty.layout.size("P")
max_inst_size = op.config.get("max_inst_size", 512)
- if ("const_bias", bias.value) in buffer_dict:
- bias_buffer, bias_init_stmt = buffer_dict[("const_bias", bias.value)]
+ if isinstance(max_inst_size, int | IntImm) and int(max_inst_size) == -1:
+ raise ValueError("Constant bias workspace allocation requires a finite
max_inst_size")
+ if bias_key in buffer_dict:
+ bias_buffer, bias_init_stmt = buffer_dict[bias_key]
old_shape = bias_buffer.ty.shape
new_shape = [max(par_size, old_shape[0]), max(max_inst_size,
old_shape[1])]
if new_shape[0] == old_shape[0] and new_shape[1] == old_shape[1]:
- return {"const_bias": ("const_bias", bias.value)}
+ return {"const_bias": bias_key}
else:
new_shape = (par_size, max_inst_size)
new_buffer = T.Var(
@@ -76,8 +81,8 @@ def alloc_const_bias_trn(
T.evaluate(T.nki.memset(new_buffer[p_loop, f_loop], bias))
T.tvm_kernel_replace_point()
- buffer_dict[("const_bias", bias.value)] = (new_buffer,
const_bias_init.body)
- return {"const_bias": ("const_bias", bias.value)}
+ buffer_dict[bias_key] = (new_buffer, const_bias_init.body)
+ return {"const_bias": bias_key}
def alloc_partial_reduce_trn(
@@ -197,7 +202,9 @@ def alloc_unary_reduce_trn(
return workspace_dict
-UnaryOpWithBiasScale.get_private_buffers_trn = alloc_const_bias_trn
+UnaryOpWithScaleBias.get_private_buffers_trn = alloc_const_bias_trn
+Sqrt.get_private_buffers_trn = alloc_const_bias_trn
+Exp.get_private_buffers_trn = alloc_const_bias_trn
ReduceOp.get_private_buffers_trn = alloc_partial_reduce_trn
Copy.get_private_buffers_trn = alloc_copy_trn
Gemm.get_private_buffers_trn = alloc_acc_psum_trn
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
index 3ca7277238..860d1b8124 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
@@ -83,7 +83,7 @@ def reduction_trn(
# Extract buffers
dst = dst_buffer_region.source
src = src_buffer_region.source
- axes = [i if i >= 0 else len(src.ty.shape) + i for i in axes]
+ axes = [int(i) if int(i) >= 0 else len(src.ty.shape) + int(i) for i in
axes]
dim_map = get_reduction_dim_map(src_buffer_region, dst_buffer_region,
axes, analyzer)
# Layout validation
diff --git a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
index 6c728c35ad..d65b794bcd 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
@@ -37,10 +37,11 @@ def unary_with_bias_scale_trn(
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
fail("requires Trainium target and thread exec_scope")
- # Extract operation arguments with defaults
- dst_buffer_region, src_buffer_region, _bias, scale = op.args
- scale = 1.0 if scale is None else scale
- _bias = 0.0 if _bias is None else _bias
+ if len(op.args) == 2:
+ dst_buffer_region, src_buffer_region = op.args
+ scale, _bias = 1.0, 0.0
+ else:
+ dst_buffer_region, src_buffer_region, scale, _bias = op.args
# Initialize analyzer and validate operation type
analyzer = init_analyzer(sctx)
@@ -84,5 +85,8 @@ from tvm.tirx.operator.tile_primitive import
register_dispatch # noqa: E402
for _op_name, _op_type in {"sqrt": MapOpType.SQRT, "exp":
MapOpType.EXP}.items():
@register_dispatch(_op_name, "trn", variant="unary_with_bias_scale",
priority=0)
+ @register_dispatch(
+ _op_name + "_with_scale_bias", "trn", variant="unary_with_bias_scale",
priority=0
+ )
def _unary_bs_dispatch(op, sctx, _ty=_op_type):
return unary_with_bias_scale_trn(op, _ty, sctx)
diff --git a/python/tvm/tirx/operator/tile_primitive/ops.py
b/python/tvm/tirx/operator/tile_primitive/ops.py
index 9b10f03374..69d0dc2760 100644
--- a/python/tvm/tirx/operator/tile_primitive/ops.py
+++ b/python/tvm/tirx/operator/tile_primitive/ops.py
@@ -58,20 +58,20 @@ class UnaryOp(TilePrimitiveCall):
return [self.output]
-class UnaryOpWithBiasScale(UnaryOp):
- """Extended unary operator with bias and scale parameters:
unary_with_bias_scale(output, input, bias, scale).
+class UnaryOpWithScaleBias(UnaryOp):
+ """Extended unary operator with bias and scale parameters:
unary_with_scale_bias(output, input, scale, bias).
- These operators support additional bias and scale parameters for more
complex operations (only on trn).
+ Both scale and bias are required operands, applied before the unary
operation.
output = unary(input * scale + bias)
""" # noqa: E501
- bias = ArgProperty(2)
- scale = ArgProperty(3)
+ scale = ArgProperty(2)
+ bias = ArgProperty(3)
@property
def srcs(self) -> list[Expr]:
"""Get the source expressions (inputs) of the operator."""
- return [self.input, self.bias, self.scale]
+ return [self.input, self.scale, self.bias]
class BinaryOp(TilePrimitiveCall):
@@ -124,15 +124,18 @@ class Zero(UnaryOp):
op = get_tirx_op("zero")
-class Sqrt(UnaryOpWithBiasScale):
- """Compute square root of all elements in src and store to dst.
-
- If bias and scale are provided: dst = sqrt(src * scale + bias)
- """
+class Sqrt(UnaryOp):
+ """Compute sqrt(src) and store to dst."""
op = get_tirx_op("sqrt")
+class SqrtWithScaleBias(UnaryOpWithScaleBias):
+ """Compute sqrt(src * scale + bias) and store to dst."""
+
+ op = get_tirx_op("sqrt_with_scale_bias")
+
+
class Fill(UnaryOp):
"""Fill dst with a scalar value."""
@@ -382,33 +385,42 @@ class Minimum(BinaryOp):
op = get_tirx_op("minimum")
-class Exp(UnaryOpWithBiasScale):
- """Compute exponential (e^x) of all elements in src and store to dst.
-
- If bias and scale are provided: dst = exp(src * scale + bias)
- """
+class Exp(UnaryOp):
+ """Compute exp(src) and store to dst."""
op = get_tirx_op("exp")
-class Exp2(UnaryOpWithBiasScale):
- """Compute base-2 exponential (2^x) of all elements in src and store to
dst.
+class ExpWithScaleBias(UnaryOpWithScaleBias):
+ """Compute exp(src * scale + bias) and store to dst."""
+
+ op = get_tirx_op("exp_with_scale_bias")
- If bias and scale are provided: dst = exp2(src * scale + bias)
- """
+
+class Exp2(UnaryOp):
+ """Compute exp2(src) and store to dst."""
op = get_tirx_op("exp2")
-class Log2(UnaryOpWithBiasScale):
- """Compute base-2 logarithm of all elements in src and store to dst.
+class Exp2WithScaleBias(UnaryOpWithScaleBias):
+ """Compute exp2(src * scale + bias) and store to dst."""
+
+ op = get_tirx_op("exp2_with_scale_bias")
- If bias and scale are provided: dst = log2(src * scale + bias)
- """
+
+class Log2(UnaryOp):
+ """Compute log2(src) and store to dst."""
op = get_tirx_op("log2")
+class Log2WithScaleBias(UnaryOpWithScaleBias):
+ """Compute log2(src * scale + bias) and store to dst."""
+
+ op = get_tirx_op("log2_with_scale_bias")
+
+
class Select(BinaryOp):
"""Select elements from src1 or src2 based on the predicate.
@@ -450,8 +462,8 @@ class BinaryReduce(TilePrimitiveCall):
class UnaryReduce(TilePrimitiveCall):
"""Combine a unary operation with a reduction operation.
- unary_reduce(unary_output, reduce_output, unary_input, unary_op,
reduce_op, bias, scale, reduce_axes)
- """ # noqa: E501
+ unary_reduce(unary_output, reduce_output, unary_input, unary_op,
reduce_op, reduce_axes)
+ """
op = get_tirx_op("unary_reduce")
@@ -460,14 +472,12 @@ class UnaryReduce(TilePrimitiveCall):
unary_input = ArgProperty(2)
unary_op = ArgProperty(3)
reduce_op = ArgProperty(4)
- bias = ArgProperty(5)
- scale = ArgProperty(6)
- reduce_axes = ArgProperty(7)
+ reduce_axes = ArgProperty(5)
@property
def srcs(self) -> list[Expr]:
"""Get the source expressions (inputs) of the operator."""
- return [self.unary_input, self.bias, self.scale]
+ return [self.unary_input]
@property
def dsts(self) -> list[Expr]:
@@ -475,6 +485,19 @@ class UnaryReduce(TilePrimitiveCall):
return [self.unary_output, self.reduce_output]
+class UnaryReduceWithScaleBias(UnaryReduce):
+ """Write unary(input * scale + bias), then reduce the unary output."""
+
+ op = get_tirx_op("unary_reduce_with_scale_bias")
+ scale = ArgProperty(5)
+ bias = ArgProperty(6)
+ reduce_axes = ArgProperty(7)
+
+ @property
+ def srcs(self) -> list[Expr]:
+ return [self.unary_input, self.scale, self.bias]
+
+
class BinaryChain(TilePrimitiveCall):
"""Chain multiple binary operations together.
diff --git a/python/tvm/tirx/script/ir_builder/tirx.py
b/python/tvm/tirx/script/ir_builder/tirx.py
index e9cec36ff6..d6051b442f 100644
--- a/python/tvm/tirx/script/ir_builder/tirx.py
+++ b/python/tvm/tirx/script/ir_builder/tirx.py
@@ -137,7 +137,7 @@ def _to_region(buffer: TensorRegion | Var):
def _wrap_elem_in_tuple(e):
- if isinstance(e, tuple | list):
+ if isinstance(e, tuple | list | tvm.ir.Array | tvm.ir.Tuple):
return e
return (e,)
@@ -184,57 +184,57 @@ def zero(
def sqrt(
dst: TensorRegion | Var,
src: TensorRegion | Var | None = None,
- bias: TensorRegion | Var | FloatImm | None = None,
- scale: FloatImm | None = None,
workspace: dict[str, Var] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
**kwargs,
):
- """Sqrt all elements in src and store to dst.
-
- dst = sqrt(src * scale + bias) (if scale or bias are provided)
-
- Parameters
- ----------
- dst : Union[TensorRegion, Var]
- The destination buffer region for sqrt result.
- When src is omitted, also used as the source (in-place).
-
- src : Union[TensorRegion, Var], optional
- The source buffer region. If omitted, dst is used (in-place).
-
- bias : Optional[Union[TensorRegion, Var, FloatImm]]
- The bias of the sqrt src. Only supported on Trn.
-
- scale : Optional[FloatImm]
- The scale of the sqrt src. Only supported on Trn.
-
- workspace : Optional[Dict[str, Var]]
- The workspace of the operator.
- """
- # Expression-form overload: ``sqrt(value)`` returns the underlying
expression.
+ """Compute sqrt(src), using dst as the source when omitted."""
from tvm import tirx as _tirx
+ if "scale" in kwargs or "bias" in kwargs:
+ raise TypeError("Use sqrt_with_scale_bias for scale and bias operands")
if not _is_buffer_or_region(dst):
return _tirx.sqrt(dst)
if src is None:
src = dst
- if workspace is None:
- workspace = {}
- config = kwargs or {}
- dst = _to_region(dst)
- src = _to_region(src)
- if bias is not None and is_tensor_var(bias):
- bias = _to_region(bias)
return f_insert(
tirx_op.Sqrt(
- dst,
- src,
- bias,
+ _to_region(dst),
+ _to_region(src),
+ workspace=workspace,
+ config=kwargs,
+ dispatch=dispatch,
+ scope=scope,
+ )
+ )
+
+
+@ScopedOp
+def sqrt_with_scale_bias(
+ dst: TensorRegion | Var,
+ src: TensorRegion | Var | None = None,
+ *,
+ scale: FloatImm,
+ bias: TensorRegion | Var | FloatImm,
+ workspace: dict[str, Var] | None = None,
+ dispatch: str | None = None,
+ scope: ExecScope | None = None,
+ **kwargs,
+):
+ """Compute sqrt(src * scale + bias), with both operands required."""
+ if src is None:
+ src = dst
+ if is_tensor_var(bias):
+ bias = _to_region(bias)
+ return f_insert(
+ tirx_op.SqrtWithScaleBias(
+ _to_region(dst),
+ _to_region(src),
scale,
+ bias,
workspace=workspace,
- config=config,
+ config=kwargs,
dispatch=dispatch,
scope=scope,
)
@@ -531,10 +531,13 @@ def _check_copy_regions_match(dst, src, config, name,
dispatch=None):
dst_bits = _payload_bits(dst)
src_bits = _payload_bits(src)
gather4 = config.get("gather4")
- if gather4 is not None:
+ if "gather4" in config:
if dispatch != "tma_explicit":
raise ValueError("copy_async: gather4 is only supported by
dispatch='tma_explicit'")
- if not isinstance(gather4, list | tuple | tvm.ir.Array) or
len(gather4) != 4:
+ if (
+ not isinstance(gather4, list | tuple | tvm.ir.Array |
tvm.ir.Tuple)
+ or len(gather4) != 4
+ ):
raise ValueError("copy_async: gather4 must contain exactly
four row coordinates")
if len(src.region) != 2:
raise ValueError("copy_async: gather4 requires a rank-2 global
source")
@@ -1163,55 +1166,26 @@ def minimum(
def exp(
dst: TensorRegion | Var,
src: TensorRegion | Var | None = None,
- bias: TensorRegion | Var | FloatImm | None = None,
- scale: FloatImm | None = None,
workspace: dict[str, Var] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
**kwargs,
):
- """Exponentiate all elements in src and store to dst.
-
- Parameters
- ----------
- dst : Union[TensorRegion, Var]
- The destination buffer region for exp result.
- When src is omitted, also used as the source (in-place).
-
- src : Union[TensorRegion, Var], optional
- The source buffer region. If omitted, dst is used (in-place).
-
- bias : Optional[Union[TensorRegion, Var, FloatImm]]
- The bias of the exp src. Only supported on Trn.
-
- scale : Optional[FloatImm]
- The scale of the exp src. Only supported on Trn.
-
- workspace : Dict[str, Var]
- The workspace of the operator.
- """
- # Expression-form overload: ``exp(value)`` returns the underlying
expression.
+ """Compute exp(src), using dst as the source when omitted."""
from tvm import tirx as _tirx
+ if "scale" in kwargs or "bias" in kwargs:
+ raise TypeError("Use exp_with_scale_bias for scale and bias operands")
if not _is_buffer_or_region(dst):
return _tirx.exp(dst)
if src is None:
src = dst
- if workspace is None:
- workspace = {}
- config = kwargs or {}
- dst = _to_region(dst)
- src = _to_region(src)
- if bias is not None and is_tensor_var(bias):
- bias = _to_region(bias)
return f_insert(
tirx_op.Exp(
- dst,
- src,
- bias,
- scale,
+ _to_region(dst),
+ _to_region(src),
workspace=workspace,
- config=config,
+ config=kwargs,
dispatch=dispatch,
scope=scope,
)
@@ -1219,58 +1193,91 @@ def exp(
@ScopedOp
-def exp2(
+def exp_with_scale_bias(
dst: TensorRegion | Var,
src: TensorRegion | Var | None = None,
- bias: TensorRegion | Var | FloatImm | None = None,
- scale: FloatImm | None = None,
+ *,
+ scale: FloatImm,
+ bias: TensorRegion | Var | FloatImm,
workspace: dict[str, Var] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
**kwargs,
):
- """Compute base-2 exponential (2^x) of all elements in src and store to
dst.
-
- Parameters
- ----------
- dst : Union[TensorRegion, Var]
- The destination buffer region for exp2 result.
- When src is omitted, also used as the source (in-place).
-
- src : Union[TensorRegion, Var], optional
- The source buffer region. If omitted, dst is used (in-place).
-
- bias : Optional[Union[TensorRegion, Var, FloatImm]]
- The bias of the exp2 src.
+ """Compute exp(src * scale + bias), with both operands required."""
+ if src is None:
+ src = dst
+ if is_tensor_var(bias):
+ bias = _to_region(bias)
+ return f_insert(
+ tirx_op.ExpWithScaleBias(
+ _to_region(dst),
+ _to_region(src),
+ scale,
+ bias,
+ workspace=workspace,
+ config=kwargs,
+ dispatch=dispatch,
+ scope=scope,
+ )
+ )
- scale : Optional[FloatImm]
- The scale of the exp2 src.
- workspace : Dict[str, Var]
- The workspace of the operator.
- """
- # Expression-form overload: ``exp2(value)`` returns the underlying
expression.
+@ScopedOp
+def exp2(
+ dst: TensorRegion | Var,
+ src: TensorRegion | Var | None = None,
+ workspace: dict[str, Var] | None = None,
+ dispatch: str | None = None,
+ scope: ExecScope | None = None,
+ **kwargs,
+):
+ """Compute exp2(src), using dst as the source when omitted."""
from tvm import tirx as _tirx
+ if "scale" in kwargs or "bias" in kwargs:
+ raise TypeError("Use exp2_with_scale_bias for scale and bias operands")
if not _is_buffer_or_region(dst):
return _tirx.exp2(dst)
if src is None:
src = dst
- if workspace is None:
- workspace = {}
- config = kwargs or {}
- dst = _to_region(dst)
- src = _to_region(src)
- if bias is not None and is_tensor_var(bias):
- bias = _to_region(bias)
return f_insert(
tirx_op.Exp2(
- dst,
- src,
- bias,
+ _to_region(dst),
+ _to_region(src),
+ workspace=workspace,
+ config=kwargs,
+ dispatch=dispatch,
+ scope=scope,
+ )
+ )
+
+
+@ScopedOp
+def exp2_with_scale_bias(
+ dst: TensorRegion | Var,
+ src: TensorRegion | Var | None = None,
+ *,
+ scale: FloatImm,
+ bias: TensorRegion | Var | FloatImm,
+ workspace: dict[str, Var] | None = None,
+ dispatch: str | None = None,
+ scope: ExecScope | None = None,
+ **kwargs,
+):
+ """Compute exp2(src * scale + bias), with both operands required."""
+ if src is None:
+ src = dst
+ if is_tensor_var(bias):
+ bias = _to_region(bias)
+ return f_insert(
+ tirx_op.Exp2WithScaleBias(
+ _to_region(dst),
+ _to_region(src),
scale,
+ bias,
workspace=workspace,
- config=config,
+ config=kwargs,
dispatch=dispatch,
scope=scope,
)
@@ -1281,36 +1288,57 @@ def exp2(
def log2(
dst: TensorRegion | Var,
src: TensorRegion | Var | None = None,
- bias: TensorRegion | Var | FloatImm | None = None,
- scale: FloatImm | None = None,
workspace: dict[str, Var] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
**kwargs,
):
- """Compute base-2 logarithm of all elements in src and store to dst."""
- # Expression-form overload: ``log2(value)`` returns the underlying
expression.
+ """Compute log2(src), using dst as the source when omitted."""
from tvm import tirx as _tirx
+ if "scale" in kwargs or "bias" in kwargs:
+ raise TypeError("Use log2_with_scale_bias for scale and bias operands")
if not _is_buffer_or_region(dst):
return _tirx.log2(dst)
if src is None:
src = dst
- if workspace is None:
- workspace = {}
- config = kwargs or {}
- dst = _to_region(dst)
- src = _to_region(src)
- if bias is not None and is_tensor_var(bias):
- bias = _to_region(bias)
return f_insert(
tirx_op.Log2(
- dst,
- src,
- bias,
+ _to_region(dst),
+ _to_region(src),
+ workspace=workspace,
+ config=kwargs,
+ dispatch=dispatch,
+ scope=scope,
+ )
+ )
+
+
+@ScopedOp
+def log2_with_scale_bias(
+ dst: TensorRegion | Var,
+ src: TensorRegion | Var | None = None,
+ *,
+ scale: FloatImm,
+ bias: TensorRegion | Var | FloatImm,
+ workspace: dict[str, Var] | None = None,
+ dispatch: str | None = None,
+ scope: ExecScope | None = None,
+ **kwargs,
+):
+ """Compute log2(src * scale + bias), with both operands required."""
+ if src is None:
+ src = dst
+ if is_tensor_var(bias):
+ bias = _to_region(bias)
+ return f_insert(
+ tirx_op.Log2WithScaleBias(
+ _to_region(dst),
+ _to_region(src),
scale,
+ bias,
workspace=workspace,
- config=config,
+ config=kwargs,
dispatch=dispatch,
scope=scope,
)
@@ -1408,83 +1436,69 @@ def unary_reduce(
unary_input: TensorRegion | Var,
unary_op: str | Op,
reduce_op: str | Op,
- bias: TensorRegion | Var | FloatImm | None = None,
- scale: FloatImm | None = None,
reduce_axes: int | tuple[int] = -1,
workspace: dict[str, Var] | None = None,
dispatch: str | None = None,
scope: ExecScope | None = None,
**kwargs,
):
- """Combine a unary operation with a reduction operation.
-
- Parameters
- ----------
- unary_output : Union[TensorRegion, Var]
- The destination buffer region for unary operation result.
-
- reduce_output : Union[TensorRegion, Var]
- The destination buffer region for reduction result.
-
- unary_input : Union[TensorRegion, Var]
- The source input for unary operation.
-
- unary_op : Union[str, Op]
- The unary operation to perform.
-
- reduce_op : Union[str, Op]
- The reduction operation to perform.
-
- bias : Optional[Union[TensorRegion, Var, FloatImm]]
- The bias to apply before unary operation.
-
- scale : Optional[FloatImm]
- The scale to apply before unary operation.
-
- reduce_axes : Union[int, Tuple[int]]
- The axes to reduce over.
-
- workspace : Dict[str, Var]
- The workspace of the operator.
-
- dispatch : str, optional
- The dispatch implementation requested for this operator.
-
- scope : ExecScope, optional
- The execution scope for this operator.
+ """Write the unary output, then its reduction."""
+ if "scale" in kwargs or "bias" in kwargs:
+ raise TypeError("Use unary_reduce_with_scale_bias for scale and bias
operands")
+ if isinstance(unary_op, str):
+ unary_op = tirx_op.get_tirx_op(unary_op)
+ if isinstance(reduce_op, str):
+ reduce_op = tirx_op.get_tirx_op(reduce_op)
+ return f_insert(
+ tirx_op.UnaryReduce(
+ _to_region(unary_output),
+ _to_region(reduce_output),
+ _to_region(unary_input),
+ unary_op,
+ reduce_op,
+ _wrap_elem_in_tuple(reduce_axes),
+ workspace=workspace,
+ config=kwargs,
+ dispatch=dispatch,
+ scope=scope,
+ )
+ )
- **kwargs
- Scheduler configuration passed as keyword arguments.
- """
- if workspace is None:
- workspace = {}
- unary_output = _to_region(unary_output)
- reduce_output = _to_region(reduce_output)
- unary_input = _to_region(unary_input)
- if bias is not None and is_tensor_var(bias):
+@ScopedOp
+def unary_reduce_with_scale_bias(
+ unary_output: TensorRegion | Var,
+ reduce_output: TensorRegion | Var,
+ unary_input: TensorRegion | Var,
+ unary_op: str | Op,
+ reduce_op: str | Op,
+ scale: FloatImm,
+ bias: TensorRegion | Var | FloatImm,
+ reduce_axes: int | tuple[int] = -1,
+ workspace: dict[str, Var] | None = None,
+ dispatch: str | None = None,
+ scope: ExecScope | None = None,
+ **kwargs,
+):
+ """Write the unary output, then its reduction. Scale and bias apply before
unary."""
+ if is_tensor_var(bias):
bias = _to_region(bias)
-
- reduce_axes = _wrap_elem_in_tuple(reduce_axes)
-
if isinstance(unary_op, str):
unary_op = tirx_op.get_tirx_op(unary_op)
if isinstance(reduce_op, str):
reduce_op = tirx_op.get_tirx_op(reduce_op)
-
- config = kwargs or {}
return f_insert(
- tirx_op.UnaryReduce(
- unary_output,
- reduce_output,
- unary_input,
+ tirx_op.UnaryReduceWithScaleBias(
+ _to_region(unary_output),
+ _to_region(reduce_output),
+ _to_region(unary_input),
unary_op,
reduce_op,
- bias,
scale,
- reduce_axes,
+ bias,
+ _wrap_elem_in_tuple(reduce_axes),
workspace=workspace,
- config=config,
+ config=kwargs,
dispatch=dispatch,
scope=scope,
)
@@ -1783,12 +1797,15 @@ __all__ = [
"cta",
"exp",
"exp2",
+ "exp2_with_scale_bias",
+ "exp_with_scale_bias",
"fdiv",
"fill",
"fma",
"gemm",
"gemm_async",
"log2",
+ "log2_with_scale_bias",
"max",
"maximum",
"memset",
@@ -1802,10 +1819,12 @@ __all__ = [
"select",
"silu",
"sqrt",
+ "sqrt_with_scale_bias",
"sub",
"sum",
"thread",
"unary_reduce",
+ "unary_reduce_with_scale_bias",
"warp",
"warpgroup",
"wg",
diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py
index a7b08d6741..208420a0a7 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -40,10 +40,18 @@ def _require_buffer_arg(op_name, arg_name, value):
def _validate_tile_call(op_name, args, kwargs):
+ if op_name in {"unary_reduce", "unary_reduce_with_scale_bias"}:
+ for index, name in enumerate(("unary_output", "reduce_output",
"unary_input")):
+ _require_buffer_arg(op_name, name, _get_arg(args, kwargs, index,
name))
+ return
dst = _get_arg(args, kwargs, 0, "dst")
_require_buffer_arg(op_name, "dst", dst)
- if op_name in {"cast", "max", "min", "permute_layout", "silu"}:
+ if op_name.endswith("_with_scale_bias"):
+ src = _get_arg(args, kwargs, 1, "src")
+ if src is not None:
+ _require_buffer_arg(op_name, "src", src)
+ elif op_name in {"cast", "max", "min", "permute_layout", "silu"}:
src = _get_arg(args, kwargs, 1, "src")
_require_buffer_arg(op_name, "src", src)
elif op_name in {"sqrt", "exp", "exp2", "log2", "reciprocal"}:
@@ -71,8 +79,11 @@ _SCOPED_TILE_OP_NAMES = [
"copy",
"copy_async",
"exp",
+ "exp_with_scale_bias",
"exp2",
+ "exp2_with_scale_bias",
"log2",
+ "log2_with_scale_bias",
"fdiv",
"fill",
"fma",
@@ -90,9 +101,11 @@ _SCOPED_TILE_OP_NAMES = [
"select",
"silu",
"sqrt",
+ "sqrt_with_scale_bias",
"sub",
"sum",
"unary_reduce",
+ "unary_reduce_with_scale_bias",
"zero",
]
diff --git a/python/tvm/tirx/tile_primitive.py
b/python/tvm/tirx/tile_primitive.py
index 48049f4a0f..a1a7edfc0a 100644
--- a/python/tvm/tirx/tile_primitive.py
+++ b/python/tvm/tirx/tile_primitive.py
@@ -31,7 +31,7 @@ from tvm.target import Target
from . import _ffi_api
from .exec_scope import ExecScope
-from .expr import FloatImm, IterVar
+from .expr import IterVar
from .stmt import Stmt
@@ -219,12 +219,6 @@ class DispatchContext(Object, Scriptable):
return self.scope_kind == "cluster"
-def normalize_const_arg(arg) -> Expr:
- if isinstance(arg, float):
- return FloatImm("float32", arg)
- return arg
-
-
@tvm_ffi.register_object("tirx.TilePrimitiveCall")
class TilePrimitiveCall(Stmt):
"""TilePrimitiveCall node.
@@ -240,8 +234,10 @@ class TilePrimitiveCall(Stmt):
workspace : Map[str, Var]
The workspace.
- config : Map[str, ObjectRef]
- The scheduler/config dictionary.
+ config : Map[str, Expr]
+ The scheduler/config dictionary. Omit unused keys; explicit None values
+ are invalid. max_inst_size=-1 means unbounded, while omission retains
+ the backend default. A present gather4 contains exactly four
coordinates.
dispatch : Optional[str]
The explicit variant name to dispatch to.
@@ -252,7 +248,7 @@ class TilePrimitiveCall(Stmt):
args: list[Expr]
workspace: dict[str, Var]
- config: dict[str, Any]
+ config: dict[str, Expr]
dispatch: str | None
scope: ExecScope
_registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {}
@@ -277,7 +273,6 @@ class TilePrimitiveCall(Stmt):
"Directly instantiating TilePrimitiveCall needs to specify the
op"
)
op = self.__class__.op
- args = list(map(normalize_const_arg, args))
self.__init_handle_by_constructor__(
_ffi_api.TilePrimitiveCall,
op,
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index fa1fe1c371..822943f94f 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -254,7 +254,8 @@ class TrainiumLayoutApplier : public
tirx::IRMutatorWithAnalyzer {
}
UnchangedOr<Stmt> Mutate_(const tirx::TilePrimitiveCallNode* op, InplaceMode
inplace_mode) final {
- auto args = op->args.Map([this](const ffi::Any& arg) { return
MutateTileArgument(arg); });
+ auto args = op->args.Map(
+ [this](const Expr& arg) { return
MutateTileArgument(arg).as_or_throw<Expr>(); });
if (args.same_as(op->args)) {
return ffi::Unchanged();
} else {
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index 6ab088995e..b1b6b7e4f9 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -833,6 +833,19 @@ PrimExpr::PrimExpr(float value) :
PrimExpr(FloatImm(PrimType::Float(32), value))
Expr ffi::TypeTraits<Expr>::ConvertFallbackValue(ffi::String value) { return
StringImm(value); }
+std::optional<Expr> ffi::TypeTraits<Expr>::TryCastFromAnyView(const TVMFFIAny*
src) {
+ if (auto value =
+ ObjectRefWithFallbackTraitsBase<Expr, PrimExpr,
ffi::String>::TryCastFromAnyView(src)) {
+ return value;
+ }
+ // Keep the recursive array conversion out of line: each element uses this
+ // trait again, while an outer Array<Expr> remains an array.
+ if (auto fields =
ffi::TypeTraits<ffi::Array<Expr>>::TryCastFromAnyView(src)) {
+ return tvm::Tuple(std::move(*fields));
+ }
+ return std::nullopt;
+}
+
namespace ffi {
PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(StrictBool value) {
diff --git a/src/relax/op/tensor/create.cc b/src/relax/op/tensor/create.cc
index d768d862e9..604b60237b 100644
--- a/src/relax/op/tensor/create.cc
+++ b/src/relax/op/tensor/create.cc
@@ -45,7 +45,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* Initialization operators */
/* relax.full */
-Expr full(ffi::Variant<Expr, ffi::Array<PrimExpr>> shape, Expr fill_value,
+// Prefer shape arrays to the generic Expr conversion of sequences to tuples.
+Expr full(ffi::Variant<ffi::Array<PrimExpr>, Expr> shape, Expr fill_value,
ffi::Optional<DLDataType> dtype) {
Expr shape_in_expr = [&]() -> Expr {
if (const auto* expr = shape.as<ExprNode>()) {
diff --git a/src/relax/op/tensor/create.h b/src/relax/op/tensor/create.h
index fea1741a54..4552878990 100644
--- a/src/relax/op/tensor/create.h
+++ b/src/relax/op/tensor/create.h
@@ -41,7 +41,7 @@ namespace relax {
* If dtype is not given, it will by default use the dtype of fill_value.
* \return The result tensor.
*/
-Expr full(ffi::Variant<Expr, ffi::Array<PrimExpr>> shape, Expr fill_value,
+Expr full(ffi::Variant<ffi::Array<PrimExpr>, Expr> shape, Expr fill_value,
ffi::Optional<DLDataType> dtype);
/*!
diff --git a/src/relax/op/tensor/manipulate.cc
b/src/relax/op/tensor/manipulate.cc
index 62b4182d4b..c9f9b7682f 100644
--- a/src/relax/op/tensor/manipulate.cc
+++ b/src/relax/op/tensor/manipulate.cc
@@ -890,7 +890,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
/* relax.reshape */
Expr ConvertNewShapeToExpr(const Expr& data,
- const ffi::Variant<Expr, ffi::Array<PrimExpr>>&
shape) {
+ const ffi::Variant<ffi::Array<PrimExpr>, Expr>&
shape) {
const ffi::ArrayObj* array;
// Treat shape expressions as constant arrays to handle special values.
if (const auto* e = shape.as<ShapeExprNode>()) {
@@ -990,7 +990,8 @@ Expr ConvertNewShapeToExpr(const Expr& data,
return ShapeExpr(array_ref);
}
-Expr reshape(Expr x, ffi::Variant<Expr, ffi::Array<PrimExpr>> shape) {
+// Prefer shape arrays to the generic Expr conversion of sequences to tuples.
+Expr reshape(Expr x, ffi::Variant<ffi::Array<PrimExpr>, Expr> shape) {
Expr shape_in_expr = ConvertNewShapeToExpr(x, shape);
static const Op op = Op::Get("relax.reshape");
return Call::Unchecked(Type::Missing(), op, {std::move(x),
std::move(shape_in_expr)}, Attrs(),
diff --git a/src/relax/op/tensor/manipulate.h b/src/relax/op/tensor/manipulate.h
index 2066c7b3b4..7525a33dfe 100644
--- a/src/relax/op/tensor/manipulate.h
+++ b/src/relax/op/tensor/manipulate.h
@@ -87,7 +87,7 @@ Expr permute_dims(Expr x, ffi::Optional<ffi::Array<int64_t>>
axes);
* It is required to be either an Array of PrimExpr, or a Shape in Relax
* \return The reshaped result.
*/
-Expr reshape(Expr x, ffi::Variant<Expr, ffi::Array<PrimExpr>> shape);
+Expr reshape(Expr x, ffi::Variant<ffi::Array<PrimExpr>, Expr> shape);
/*!
* \brief Split input tensor along axis by sections or indices.
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index ae1b53fac8..530dde0e2e 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -213,27 +213,11 @@ ffi::Optional<VisitInterrupt>
StmtExprVisitor::Visit_(const ScopeIdDefStmtNode*
}
ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const
TilePrimitiveCallNode* op) {
- std::function<ffi::Optional<VisitInterrupt>(const ffi::Any&)> fvisit;
- fvisit = [this, &fvisit](const ffi::Any& e) -> ffi::Optional<VisitInterrupt>
{
- if (e == nullptr) return std::nullopt;
- if (auto buffer_region = e.as<TensorRegion>()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(buffer_region.value()));
- } else if (auto expr = e.as<Expr>()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(expr.value()));
- } else if (auto stmt = e.as<Stmt>()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(stmt.value()));
- } else if (auto array = e.as<ffi::Array<ffi::Any>>()) {
- for (const ffi::Any& item : array.value()) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit(item));
- }
- }
- return std::nullopt;
- };
- for (const ffi::Any& arg : op->args) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit(arg));
+ for (const Expr& arg : op->args) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(arg));
}
for (const auto& [key, value] : op->config) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(fvisit(value));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(value));
}
return std::nullopt;
}
@@ -460,33 +444,6 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const
SeqStmtNode* op, InplaceMode in
});
}
-namespace {
-
-template <typename T, typename F>
-UnchangedOr<ffi::Array<T>> MutateTileArray(const ffi::ArrayObj* values,
InplaceMode mode,
- F fmutate) {
- // Borrow both the owning container and its elements throughout recursion.
- if (!values->unique()) mode = InplaceMode::kDisallow;
- std::vector<std::pair<size_t, T>> replacements;
- for (size_t i = 0; i < values->size(); ++i) {
- UnchangedOr<ffi::Any> result = fmutate(ffi::AnyView((*values)[i]), mode);
- if (!result.UnchangedOrSameAs((*values)[i])) {
- replacements.emplace_back(i, std::move(result).ValueUnchecked().template
as_or_throw<T>());
- }
- }
- if (replacements.empty()) return ffi::Unchanged();
- if (mode == InplaceMode::kAllow) {
- auto* writable = const_cast<ffi::ArrayObj*>(values);
- for (auto& [i, value] : replacements) writable->SetItem(i,
std::move(value));
- return ffi::Unchanged();
- }
- ffi::Array<T>
result(ffi::GetObjectPtr<ffi::ArrayObj>(const_cast<ffi::ArrayObj*>(values)));
- for (auto& [i, value] : replacements) result.Set(i, std::move(value));
- return result;
-}
-
-} // namespace
-
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const ScopeIdDefStmtNode* op,
InplaceMode inplace_mode) {
// The definition owns both optional arrays; it is skipped by this semantic
hook.
InplaceMode def_mode = op->def.unique() ? inplace_mode :
InplaceMode::kDisallow;
@@ -510,39 +467,9 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const
ScopeIdDefStmtNode* op, Inplace
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const TilePrimitiveCallNode* op,
InplaceMode inplace_mode) {
- std::function<UnchangedOr<ffi::Any>(ffi::AnyView, InplaceMode)> mutate_arg;
- mutate_arg = [&](ffi::AnyView value, InplaceMode mode) ->
UnchangedOr<ffi::Any> {
- if (value.as<TensorRegionNode>()) {
- return Mutate(value, mode);
- }
- if (value.as<Expr>() || value.as<StmtNode>()) return Mutate(value, mode);
- if (const auto* array = value.as<ffi::ArrayObj>()) {
- return MutateTileArray<ffi::Any>(array, mode, mutate_arg);
- }
- return ffi::Unchanged();
- };
- auto args = MutateTileArray<ffi::Any>(op->args.GetArrayObj(), inplace_mode,
mutate_arg);
- // A config map is another owning container on the path to its values.
- auto config_mode = op->config.unique() ? inplace_mode :
InplaceMode::kDisallow;
- UnchangedOr<ffi::Map<ffi::String, ffi::Any>> config = ffi::Unchanged();
- std::vector<std::pair<ffi::String, ffi::Any>> replacements;
- for (const auto& [key, value] : *static_cast<const
ffi::MapObj*>(op->config.get())) {
- auto result = mutate_arg(value, config_mode);
- if (!result.UnchangedOrSameAs(value)) {
- replacements.emplace_back(key.as_or_throw<ffi::String>(),
std::move(result).ValueUnchecked());
- }
- }
- if (!replacements.empty()) {
- if (config_mode == InplaceMode::kAllow) {
- for (auto& [key, value] : replacements)
- const_cast<ffi::MapObj*>(static_cast<const
ffi::MapObj*>(op->config.get()))->at(key) =
- std::move(value);
- } else {
- ffi::Map<ffi::String, ffi::Any> replacement = op->config;
- for (auto& [key, value] : replacements) replacement.Set(key,
std::move(value));
- config = std::move(replacement);
- }
- }
+ auto args = Mutate(op->args,
inplace_mode).as_or_throw<UnchangedOr<ffi::Array<Expr>>>();
+ auto config =
+ Mutate(op->config,
inplace_mode).as_or_throw<UnchangedOr<ffi::Map<ffi::String, Expr>>>();
if (args.UnchangedOrSameAs(op->args) &&
config.UnchangedOrSameAs(op->config)) {
return ffi::Unchanged();
}
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index 54cd11dbe8..6f7c072f11 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -251,21 +251,22 @@ void TIRVisitorWithPath::Dispatch_(const EvaluateNode*
op, AccessPath path) {
}
void TIRVisitorWithPath::Dispatch_(const tirx::TilePrimitiveCallNode* op,
AccessPath path) {
- for (size_t i = 0; i < op->args.size(); i++) {
- if (op->args[i] == nullptr) {
- continue;
- }
- if (auto buf_region = op->args[i].as<TensorRegion>()) {
- Visit(buf_region.value(), path->Attr("args")->ArrayItem(i));
- } else if (auto lambda = op->args[i].as<LambdaExpr>()) {
- Visit(lambda.value(), path->Attr("args")->ArrayItem(i));
- } else if (auto expr = op->args[i].as<PrimExpr>()) {
- Visit(expr.value(), path->Attr("args")->ArrayItem(i));
- } else if (auto stmt = op->args[i].as<Stmt>()) {
- Visit(stmt.value(), path->Attr("args")->ArrayItem(i));
- } else if (auto buf = op->args[i].as<TensorVar>()) {
- VisitBufferUse(buf.value(), path->Attr("args")->ArrayItem(i));
+ std::function<void(const Expr&, AccessPath)> visit = [&](const Expr& expr,
AccessPath path) {
+ if (auto buffer = expr.as<TensorVar>()) {
+ VisitBufferUse(buffer.value(), path);
+ } else if (const auto* tuple = expr.as<TupleNode>()) {
+ for (size_t i = 0; i < tuple->fields.size(); ++i) {
+ visit(tuple->fields[i], path->Attr("fields")->ArrayItem(i));
+ }
+ } else if (!expr.as<OpNode>()) {
+ Visit(expr, path);
}
+ };
+ for (size_t i = 0; i < op->args.size(); ++i) {
+ visit(op->args[i], path->Attr("args")->ArrayItem(i));
+ }
+ for (const auto& [key, value] : op->config) {
+ visit(value, path->Attr("config")->MapItem(key));
}
}
diff --git a/src/tirx/ir/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc
index 09a7227bcb..39e081aefb 100644
--- a/src/tirx/ir/tirx_stmt.cc
+++ b/src/tirx/ir/tirx_stmt.cc
@@ -35,15 +35,26 @@ namespace tvm {
namespace tirx {
// TilePrimitiveCall
-TilePrimitiveCall::TilePrimitiveCall(tvm::Op op, ffi::Array<ffi::Any> args,
+TilePrimitiveCall::TilePrimitiveCall(tvm::Op op, ffi::Array<Expr> args,
ffi::Map<ffi::String, TensorVar>
workspace,
- ffi::Map<ffi::String, ffi::Any> config,
+ ffi::Map<ffi::String, Expr> config,
ffi::Optional<ffi::String> dispatch,
ExecScope scope)
: Stmt(ffi::UnsafeInit{}) {
TVM_FFI_CHECK(op.defined(), ValueError) << "TilePrimitiveCall expects a
defined operator";
static const auto& category_map =
Op::GetAttrMap<TIRxOpCategory>("TIRxOpCategory");
TVM_FFI_ICHECK(category_map.get(op, ffi::String("")) == "tile_primitive")
<< "Only tile primitive ops can be used in tirx::TilePrimitiveCall";
+ TVM_FFI_CHECK_GE(args.size(), op->args_info.size(), ValueError)
+ << op->name << " requires " << op->args_info.size() << " operands";
+ if (!op->var_args_info.has_value()) {
+ TVM_FFI_CHECK_EQ(args.size(), op->args_info.size(), ValueError)
+ << op->name << " expects " << op->args_info.size() << " operands";
+ }
+ if (auto gather = config.Get("gather4")) {
+ auto tuple = gather.value().as<tvm::Tuple>();
+ TVM_FFI_CHECK(tuple.has_value() && tuple.value()->fields.size() == 4,
ValueError)
+ << "gather4 must contain exactly four row coordinates";
+ }
ffi::StructuralVisit(args,
[](const TensorRegionNode* region,
ffi::StructuralVisitorObj*) ->
ffi::Optional<ffi::VisitInterrupt> {
@@ -78,11 +89,11 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TilePrimitiveCallMutate
const TilePrimitiveCallNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TilePrimitiveCallNode>(
value);
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<ffi::Any>>,
mapped_args,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Expr>>,
mapped_args,
mutator->MutateExpected(self->args));
using WorkspaceMap = ffi::Map<ffi::String, TensorVar>;
- using ConfigMap = ffi::Map<ffi::String, ffi::Any>;
+ using ConfigMap = ffi::Map<ffi::String, Expr>;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<WorkspaceMap>,
mapped_workspace,
mutator->MutateExpected(self->workspace));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ConfigMap>, mapped_config,
@@ -106,11 +117,11 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TilePrimitiveCallMaybeI
TilePrimitiveCallNode* self = const_cast<TilePrimitiveCallNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TilePrimitiveCallNode>(
value));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<ffi::Any>>,
mapped_args,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Expr>>,
mapped_args,
mutator->MutateExpected(self->args,
ffi::InplaceMode::kAllow));
using WorkspaceMap = ffi::Map<ffi::String, TensorVar>;
- using ConfigMap = ffi::Map<ffi::String, ffi::Any>;
+ using ConfigMap = ffi::Map<ffi::String, Expr>;
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
ffi::UnchangedOr<WorkspaceMap>, mapped_workspace,
mutator->MutateExpected(self->workspace, ffi::InplaceMode::kAllow));
@@ -136,9 +147,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TilePrimitiveCallNode::RegisterReflection();
refl::GlobalDef().def(
"tirx.TilePrimitiveCall",
- [](tvm::Op op, ffi::Array<ffi::Any> args, ffi::Map<ffi::String,
TensorVar> workspace,
- ffi::Map<ffi::String, ffi::Any> config, ffi::Optional<ffi::String>
dispatch,
- ExecScope scope) {
+ [](tvm::Op op, ffi::Array<Expr> args, ffi::Map<ffi::String, TensorVar>
workspace,
+ ffi::Map<ffi::String, Expr> config, ffi::Optional<ffi::String>
dispatch, ExecScope scope) {
return TilePrimitiveCall(op, args, workspace, config, dispatch, scope);
});
refl::TypeAttrDef<TilePrimitiveCallNode>()
diff --git a/src/tirx/op/tirx.cc b/src/tirx/op/tirx.cc
index 2eb26db1ca..1f07dfd3c7 100644
--- a/src/tirx/op/tirx.cc
+++ b/src/tirx/op/tirx.cc
@@ -138,9 +138,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TVM_DEFINE_CACHED_OP_GETTER(zero, "tirx.tile.zero")
TVM_DEFINE_CACHED_OP_GETTER(sqrt, "tirx.tile.sqrt")
+TVM_DEFINE_CACHED_OP_GETTER(sqrt_with_scale_bias,
"tirx.tile.sqrt_with_scale_bias")
TVM_DEFINE_CACHED_OP_GETTER(exp, "tirx.tile.exp")
+TVM_DEFINE_CACHED_OP_GETTER(exp_with_scale_bias,
"tirx.tile.exp_with_scale_bias")
TVM_DEFINE_CACHED_OP_GETTER(exp2, "tirx.tile.exp2")
+TVM_DEFINE_CACHED_OP_GETTER(exp2_with_scale_bias,
"tirx.tile.exp2_with_scale_bias")
TVM_DEFINE_CACHED_OP_GETTER(log2, "tirx.tile.log2")
+TVM_DEFINE_CACHED_OP_GETTER(log2_with_scale_bias,
"tirx.tile.log2_with_scale_bias")
TVM_DEFINE_CACHED_OP_GETTER(add, "tirx.tile.add")
TVM_DEFINE_CACHED_OP_GETTER(sub, "tirx.tile.sub")
TVM_DEFINE_CACHED_OP_GETTER(mul, "tirx.tile.mul")
@@ -158,6 +162,7 @@ TVM_DEFINE_CACHED_OP_GETTER(memset, "tirx.tile.memset")
TVM_DEFINE_CACHED_OP_GETTER(reduce_negate, "tirx.tile.reduce_negate")
TVM_DEFINE_CACHED_OP_GETTER(binary_reduce, "tirx.tile.binary_reduce")
TVM_DEFINE_CACHED_OP_GETTER(unary_reduce, "tirx.tile.unary_reduce")
+TVM_DEFINE_CACHED_OP_GETTER(unary_reduce_with_scale_bias,
"tirx.tile.unary_reduce_with_scale_bias")
TVM_DEFINE_CACHED_OP_GETTER(binary_chain, "tirx.tile.binary_chain")
TVM_DEFINE_CACHED_OP_GETTER(select, "tirx.tile.select")
TVM_DEFINE_CACHED_OP_GETTER(cast, "tirx.tile.cast")
@@ -176,29 +181,53 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
OpDef("tirx.tile.sqrt")
- .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."),
- sig::var_args("args"))
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tile.sqrt"))
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
- OpDef("tirx.tile.exp")
+ OpDef("tirx.tile.sqrt_with_scale_bias")
.signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."),
- sig::var_args("args"))
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The
bias."))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.tile.sqrt_with_scale_bias"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+
+ OpDef("tirx.tile.exp")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tile.exp"))
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
- OpDef("tirx.tile.exp2")
+ OpDef("tirx.tile.exp_with_scale_bias")
.signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."),
- sig::var_args("args"))
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The
bias."))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.tile.exp_with_scale_bias"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+
+ OpDef("tirx.tile.exp2")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tile.exp2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
- OpDef("tirx.tile.log2")
+ OpDef("tirx.tile.exp2_with_scale_bias")
.signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."),
- sig::var_args("args"))
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The
bias."))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.tile.exp2_with_scale_bias"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+
+ OpDef("tirx.tile.log2")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tile.log2"))
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+ OpDef("tirx.tile.log2_with_scale_bias")
+ .signature(sig::arg("dst", "The destination."), sig::arg("src", "The
source."),
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The
bias."))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.tile.log2_with_scale_bias"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+
OpDef("tirx.tile.add")
.signature(sig::arg("dst", "The destination."), sig::arg("src1", "The
first source."),
sig::arg("src2", "The second source."))
@@ -306,12 +335,23 @@ TVM_FFI_STATIC_INIT_BLOCK() {
sig::arg("reduce_output", "The reduction output."),
sig::arg("unary_input", "The unary operation input."),
sig::arg("unary_op", "The unary operation."),
- sig::arg("reduce_op", "The reduction operation."),
sig::arg("bias", "The bias."),
- sig::arg("scale", "The scale factor."),
+ sig::arg("reduce_op", "The reduction operation."),
sig::arg("reduce_axes", "The reduction axes."))
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tile.unary_reduce"))
.set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+ OpDef("tirx.tile.unary_reduce_with_scale_bias")
+ .signature(sig::arg("unary_output", "The unary operation output."),
+ sig::arg("reduce_output", "The reduction output."),
+ sig::arg("unary_input", "The unary operation input."),
+ sig::arg("unary_op", "The unary operation."),
+ sig::arg("reduce_op", "The reduction operation."),
+ sig::arg("scale", "The scale factor."), sig::arg("bias", "The
bias."),
+ sig::arg("reduce_axes", "The reduction axes."))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.tile.unary_reduce_with_scale_bias"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory",
ffi::String("tile_primitive"));
+
OpDef("tirx.tile.binary_chain")
.signature(sig::arg("output", "The output."), sig::arg("data", "The
input data."),
sig::arg("operand0", "The first operand."),
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 352fc1d69e..08695fb90a 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -25,6 +25,7 @@
#include <tvm/tirx/tile_primitive.h>
#include <algorithm>
+#include <cmath>
#include <optional>
#include <string>
#include <utility>
@@ -78,7 +79,6 @@ ffi::Optional<ExprDoc>
TilePrimitiveCallDocTranslate(DocTranslatorObj* d, ffi::A
}
ffi::Array<Doc> args;
size_t n = stmt->args.size();
- while (n && stmt->args[n - 1].type_index() == ffi::TypeIndex::kTVMFFINone)
--n;
if (n == 2 &&
(stmt->op->name == "tirx.tile.exp2" || stmt->op->name ==
"tirx.tile.sqrt" ||
stmt->op->name == "tirx.tile.reciprocal") &&
@@ -115,7 +115,7 @@ ffi::Optional<ExprDoc>
TilePrimitiveCallDocTranslate(DocTranslatorObj* d, ffi::A
args.push_back(AnyValue(d, stmt->args[i]));
}
}
- auto dict = [&](const auto& source) -> ffi::Optional<DictDoc> {
+ auto dict = [&](const auto& source, bool config = false) ->
ffi::Optional<DictDoc> {
if (source.empty()) return std::nullopt;
std::vector<std::pair<ffi::String, ffi::Any>> sorted;
for (const auto& [key, value] : source) sorted.emplace_back(key, value);
@@ -125,6 +125,31 @@ ffi::Optional<ExprDoc>
TilePrimitiveCallDocTranslate(DocTranslatorObj* d, ffi::A
ffi::Array<ExprDoc> values;
for (const auto& [key, value] : sorted) {
keys.push_back(LiteralDoc::Str(key, std::nullopt));
+ if (config) {
+ if (const auto* string = value.template as<StringImmNode>()) {
+ values.push_back(LiteralDoc::Str(string->value, std::nullopt));
+ continue;
+ }
+ if (const auto* integer = value.template as<IntImmNode>()) {
+ PrimType type = integer->ty.template as_or_throw<PrimType>();
+ if (type == PrimType::Bool()) {
+ values.push_back(LiteralDoc::Boolean(integer->value != 0,
std::nullopt));
+ continue;
+ }
+ int bits = integer->value >= INT32_MIN && integer->value <=
INT32_MAX ? 32 : 64;
+ if (type == PrimType::Int(bits)) {
+ values.push_back(LiteralDoc::Int(ffi::GetRef<IntImm>(integer),
std::nullopt));
+ continue;
+ }
+ }
+ if (const auto* floating = value.template as<FloatImmNode>()) {
+ if (floating->ty.template as_or_throw<PrimType>() ==
PrimType::Float(32) &&
+ std::isfinite(floating->value)) {
+ values.push_back(LiteralDoc::Float(floating->value, std::nullopt));
+ continue;
+ }
+ }
+ }
values.push_back(AnyValue(d, value));
}
return DictDoc(keys, values);
@@ -133,8 +158,21 @@ ffi::Optional<ExprDoc>
TilePrimitiveCallDocTranslate(DocTranslatorObj* d, ffi::A
if (stmt->dispatch.has_value()) {
dispatch = LiteralDoc::Str(stmt->dispatch.value(), std::nullopt);
}
+ auto keywords = dict(stmt->config, true);
+ if (name == "sqrt_with_scale_bias" || name == "exp_with_scale_bias" ||
+ name == "exp2_with_scale_bias" || name == "log2_with_scale_bias") {
+ ffi::Array<ExprDoc> keys{LiteralDoc::Str("scale", std::nullopt),
+ LiteralDoc::Str("bias", std::nullopt)};
+ ffi::Array<ExprDoc> values{args[2].as_or_throw<ExprDoc>(),
args[3].as_or_throw<ExprDoc>()};
+ if (keywords.has_value()) {
+ for (const auto& key : keywords.value()->keys) keys.push_back(key);
+ for (const auto& value : keywords.value()->values)
values.push_back(value);
+ }
+ keywords = DictDoc(keys, values);
+ args = {args[0], args[1]};
+ }
d->Emit(OpCallDoc(NamespaceDoc("tirx")->Attr(scope)->Attr(name), args,
dict(stmt->workspace),
- dict(stmt->config), dispatch),
+ keywords, dispatch),
ffi::GetRef<ffi::ObjectRef>(stmt));
return std::nullopt;
}
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc
b/src/tirx/transform/lower_tirx_cleanup.cc
index 36dfdbeb85..308d7cdd9c 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -276,8 +276,8 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
}
UnchangedOr<Stmt> Mutate_(const tirx::TilePrimitiveCallNode* op, InplaceMode
inplace_mode) final {
- ffi::Array<ffi::Any> args = op->args;
- args.MutateByApply([this](ffi::Any arg) -> ffi::Any { return
VisitAny(arg); });
+ ffi::Array<Expr> args = op->args;
+ args.MutateByApply([this](const Expr& arg) { return
VisitAny(arg).as_or_throw<Expr>(); });
if (args.same_as(op->args)) {
return ffi::Unchanged();
} else {
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
index 0671c8a76f..3cf2546e70 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
@@ -263,33 +263,33 @@ def test_unary_op_shared_with_bias_scale(input, op_type,
bias_type, src_dtype, d
T.cuda.cta_sync()
if bias_type == "const":
if op_type == "sqrt":
- Tx.cta.sqrt(
+ Tx.cta.sqrt_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- const_bias,
- scale,
+ scale=scale,
+ bias=const_bias,
)
elif op_type == "exp":
- Tx.cta.exp(
+ Tx.cta.exp_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- const_bias,
- scale,
+ scale=scale,
+ bias=const_bias,
)
elif bias_type == "region":
if op_type == "sqrt":
- Tx.cta.sqrt(
+ Tx.cta.sqrt_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- bias_smem[tuple(map_slice_a)],
- scale,
+ scale=scale,
+ bias=bias_smem[tuple(map_slice_a)],
)
elif op_type == "exp":
- Tx.cta.exp(
+ Tx.cta.exp_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- bias_smem[tuple(map_slice_a)],
- scale,
+ scale=scale,
+ bias=bias_smem[tuple(map_slice_a)],
)
T.cuda.cta_sync()
Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
@@ -312,33 +312,33 @@ def test_unary_op_shared_with_bias_scale(input, op_type,
bias_type, src_dtype, d
T.cuda.cta_sync()
if bias_type == "const":
if op_type == "sqrt":
- Tx.cta.sqrt(
+ Tx.cta.sqrt_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- const_bias,
- scale,
+ scale=scale,
+ bias=const_bias,
)
elif op_type == "exp":
- Tx.cta.exp(
+ Tx.cta.exp_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- const_bias,
- scale,
+ scale=scale,
+ bias=const_bias,
)
elif bias_type == "region":
if op_type == "sqrt":
- Tx.cta.sqrt(
+ Tx.cta.sqrt_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- bias_smem[tuple(map_slice_a)],
- scale,
+ scale=scale,
+ bias=bias_smem[tuple(map_slice_a)],
)
elif op_type == "exp":
- Tx.cta.exp(
+ Tx.cta.exp_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
- bias_smem[tuple(map_slice_a)],
- scale,
+ scale=scale,
+ bias=bias_smem[tuple(map_slice_a)],
)
T.cuda.cta_sync()
Tx.cta.copy(B[tuple(map_slice_res)], B_smem[tuple(map_slice_res)])
@@ -657,14 +657,14 @@ def test_unary_op_local_with_bias_scale(input, op_type,
bias_type, src_dtype, ds
bias_view = bias_local.view(*bias_shape, layout=acc_layout)
if bias_type == "const":
if op_type == "sqrt":
- Tx.warp.sqrt(res_view, acc_view, const_bias, scale)
+ Tx.warp.sqrt_with_scale_bias(res_view, acc_view, scale=scale,
bias=const_bias)
elif op_type == "exp":
- Tx.warp.exp(res_view, acc_view, const_bias, scale)
+ Tx.warp.exp_with_scale_bias(res_view, acc_view, scale=scale,
bias=const_bias)
elif bias_type == "region":
if op_type == "sqrt":
- Tx.warp.sqrt(res_view, acc_view, bias_view, scale)
+ Tx.warp.sqrt_with_scale_bias(res_view, acc_view, scale=scale,
bias=bias_view)
elif op_type == "exp":
- Tx.warp.exp(res_view, acc_view, bias_view, scale)
+ Tx.warp.exp_with_scale_bias(res_view, acc_view, scale=scale,
bias=bias_view)
# write res into B
for i in T.serial(NUM_COL // 8):
diff --git
a/tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py
index c292f06169..5aa3ffd39f 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_compose_op_trn.py
@@ -241,7 +241,7 @@ def test_activation_reduce_with_bias_scale():
C = T.alloc_tensor(C_shape, dtype="float32", scope="trn.sbuf",
layout=C_layout)
bias = T.alloc_tensor(bias_shape, dtype="float32", scope="trn.sbuf",
layout=bias_layout)
for i in range(2):
- Tx.unary_reduce(B, C, A[i*16:i*16+16], "sqrt", "sum",
reduce_axes=1, bias=bias, scale=2.0) # noqa: E501
+ Tx.unary_reduce_with_scale_bias(B, C, A[i*16:i*16+16], "sqrt",
"sum", reduce_axes=1, bias=bias, scale=2.0) # noqa: E501
@T.prim_func
def expected():
@@ -746,7 +746,7 @@ def test_unary_reduce_complex():
qk = T.alloc_tensor((2, 128, 8192), scope="trn.sbuf", layout="FPF")
running_max = T.alloc_tensor((16384, 1), dtype="float32",
scope="trn.sbuf", layout="PF")
for i in range(4):
- Tx.unary_reduce(p[0:128, 0:8192], rowsum_p[i % 2, 0:128, 0], qk[i
% 2, 0:128, 0:8192], "exp", "sum", bias=running_max[i * 128:i * 128 + 128, 0])
# noqa: E501
+ Tx.unary_reduce_with_scale_bias(p[0:128, 0:8192], rowsum_p[i % 2,
0:128, 0], qk[i % 2, 0:128, 0:8192], "exp", "sum", bias=running_max[i * 128:i *
128 + 128, 0], scale=1.0) # noqa: E501
@T.prim_func
def expected():
diff --git
a/tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py
index 21082e0a6e..f73a2cf73c 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_private_alloc_trn.py
@@ -97,7 +97,7 @@ def test_unary_with_bias_scale():
T.device_entry()
A_sbuf = T.alloc_tensor(src_shape, "float32", scope="trn.sbuf",
layout=src_layout)
C_sbuf = T.alloc_tensor(dst_shape, "float32", scope="trn.sbuf",
layout=dst_layout)
- Tx.exp(C_sbuf, A_sbuf, bias=bias, scale=scale)
+ Tx.exp_with_scale_bias(C_sbuf, A_sbuf, bias=bias, scale=scale)
@T.prim_func
def expected():
@@ -112,7 +112,7 @@ def test_unary_with_bias_scale():
layout=T.TileLayout(T.S[(128, 4096) : (1@P,
1@F)]))
C_sbuf = T.alloc_tensor((512, 1024), scope="trn.sbuf",
layout=T.TileLayout(T.S[(128, 4096) : (1@P,
1@F)]))
- Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(1.0),
T.float32(2.0), workspace={"const_bias": const_bias}) # noqa: E501
+ Tx.exp_with_scale_bias(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024],
scale=T.float32(2.0), bias=T.float32(1.0), workspace={"const_bias":
const_bias}) # noqa: E501
# fmt: on
with target:
mod = tvm.IRModule({"main": unary})
@@ -264,7 +264,7 @@ def test_activation_reduce_two_stage():
C = T.alloc_tensor((1, 128), scope="trn.sbuf",
layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
for i in range(2):
- Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i *
16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1],
workspace={"const_bias": const_bias, "partial_reduce": partial_reduce}) #
noqa: E501
+ Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i *
16 + 16, 0:512, 0:128], "sqrt", "sum", [0, 1], workspace={"const_bias":
const_bias, "partial_reduce": partial_reduce}) # noqa: E501
# fmt: on
with target:
mod = tvm.IRModule({"main": activation_reduce})
@@ -308,7 +308,7 @@ def test_partial_workspace_specify():
C = T.alloc_tensor((1, 128), scope="trn.sbuf",
layout=T.TileLayout(T.S[(1, 128) : (1@F, 1@P)]))
for i in range(2):
- Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i *
16 + 16, 0:512, 0:128], "sqrt", "sum", None, None, [0, 1],
workspace={"const_bias": const_bias, "partial_reduce": partial_reduce}) #
noqa: E501
+ Tx.unary_reduce(B[0:16, 0:512, 0:128], C[0, 0:128], A[i * 16:i *
16 + 16, 0:512, 0:128], "sqrt", "sum", [0, 1], workspace={"const_bias":
const_bias, "partial_reduce": partial_reduce}) # noqa: E501
# fmt: on
with target:
mod = tvm.IRModule({"main": activation_reduce})
@@ -329,7 +329,7 @@ def test_workspace_reuse():
T.device_entry()
A_sbuf = T.alloc_tensor(src_shape, "float32", scope="trn.sbuf",
layout=src_layout)
C_sbuf = T.alloc_tensor(dst_shape, "float32", scope="trn.sbuf",
layout=dst_layout)
- Tx.exp(C_sbuf, A_sbuf, bias=0.0, scale=scale, max_inst_size=1024)
+ Tx.exp_with_scale_bias(C_sbuf, A_sbuf, bias=0.0, scale=scale,
max_inst_size=1024)
Tx.exp(C_sbuf, C_sbuf)
@T.prim_func
@@ -345,8 +345,8 @@ def test_workspace_reuse():
layout=T.TileLayout(T.S[(128, 4096) : (1 @ P,
1 @ F)]))
C_sbuf = T.alloc_tensor((512, 1024), scope="trn.sbuf",
layout=T.TileLayout(T.S[(128, 4096) : (1 @ P,
1 @ F)]))
- Tx.exp(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024], T.float32(0.0),
T.float32(2.0), workspace={"const_bias": const_bias}, max_inst_size=1024) #
noqa: E501
- Tx.exp(C_sbuf[0:512, 0:1024], C_sbuf[0:512, 0:1024], None, None,
workspace={"const_bias": const_bias}) # noqa: E501
+ Tx.exp_with_scale_bias(C_sbuf[0:512, 0:1024], A_sbuf[0:512, 0:1024],
scale=T.float32(2.0), bias=T.float32(0.0), workspace={"const_bias":
const_bias}, max_inst_size=1024) # noqa: E501
+ Tx.exp(C_sbuf[0:512, 0:1024], C_sbuf[0:512, 0:1024],
workspace={"const_bias": const_bias})
# fmt: on
diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py
index e3287afd3d..491ef43197 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_unary_trn.py
@@ -173,7 +173,7 @@ def test_unary_with_bias_scale(op_type):
bias_shape = [512, 1]
bias_layout = TileLayout(S[(128, 4) : (1 @ P, 1 @ F)])
scale = T.float32(2.0)
- tx_func = Tx_func_map[op_type]
+ tx_func = getattr(Tx, op_type + "_with_scale_bias")
# fmt: off
@T.prim_func
@@ -210,7 +210,7 @@ def test_unary_with_bias_scale_2(op_type):
dst_layout = src_layout
bias = T.float32(1.0)
scale = T.float32(2.0)
- tx_func = Tx_func_map[op_type]
+ tx_func = getattr(Tx, op_type + "_with_scale_bias")
# fmt: off
@T.prim_func
@@ -261,7 +261,7 @@ def test_unary_with_guard():
C_sbuf = T.alloc_tensor(dst_shape, "float32", scope="trn.sbuf",
layout=dst_layout)
for i in range(4):
for j in range(4):
- Tx.sqrt(C_sbuf[0: (i+1) * 128, 0: (j+1)*256], A_sbuf[0: (i+1)
* 128, 0: (j+1)*256], bias=B_sbuf[0: (i+1) * 128, 0], scale=scale) # noqa: E501
+ Tx.sqrt_with_scale_bias(C_sbuf[0: (i+1) * 128, 0: (j+1)*256],
A_sbuf[0: (i+1) * 128, 0: (j+1)*256], bias=B_sbuf[0: (i+1) * 128, 0],
scale=scale) # noqa: E501
@T.prim_func
def expected():
diff --git a/tests/python/tirx/script/test_tirx_script_gpu.py
b/tests/python/tirx/script/test_tirx_script_gpu.py
index 4a582aba15..c22872d301 100644
--- a/tests/python/tirx/script/test_tirx_script_gpu.py
+++ b/tests/python/tirx/script/test_tirx_script_gpu.py
@@ -315,7 +315,7 @@ def test_workspace_default_none():
assert len(op_br.workspace) == 0
op_ur = tirx_op.UnaryReduce(
- B, C, A, tirx_op.get_tirx_op("sqrt"), tirx_op.get_tirx_op("sum"),
None, None, (-1,)
+ B, C, A, tirx_op.get_tirx_op("sqrt"), tirx_op.get_tirx_op("sum"), (-1,)
)
assert len(op_ur.workspace) == 0
diff --git a/tests/python/tirx/test_hint.py b/tests/python/tirx/test_hint.py
index 46b4322e9c..5b5a352df6 100644
--- a/tests/python/tirx/test_hint.py
+++ b/tests/python/tirx/test_hint.py
@@ -43,7 +43,7 @@ def test_hint_keyword_arg_on_tx_op():
config={"hint": "3-input ptx"},
)
assert "hint" in op_call.config
- assert str(op_call.config["hint"]) == "3-input ptx"
+ assert op_call.config["hint"].value == "3-input ptx"
def test_hint_keyword_arg_on_tx_op_roundtrip():
