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 1be3bc9bf0 [REFACTOR][TRN] Use RegionStmt for tensorized NKI
instructions (#20560)
1be3bc9bf0 is described below
commit 1be3bc9bf0ccd503c6175ef7ca57b81c081624ac
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 12:18:46 2026 -0400
[REFACTOR][TRN] Use RegionStmt for tensorized NKI instructions (#20560)
---
include/tvm/tirx/stmt.h | 1 -
python/tvm/backend/trn/script.py | 9 +++
.../backend/trn/tile_primitive/binary/default.py | 2 +-
.../trn/tile_primitive/compose_op/binary_chain.py | 2 +-
.../trn/tile_primitive/compose_op/binary_reduce.py | 6 +-
.../trn/tile_primitive/compose_op/unary_reduce.py | 6 +-
.../tvm/backend/trn/tile_primitive/copy/default.py | 10 ++--
.../tvm/backend/trn/tile_primitive/gemm/default.py | 4 +-
.../backend/trn/tile_primitive/private_alloc.py | 4 +-
.../backend/trn/tile_primitive/reduction/utils.py | 6 +-
.../backend/trn/tile_primitive/select/default.py | 2 +-
.../tvm/backend/trn/tile_primitive/unary/utils.py | 4 +-
src/backend/trn/codegen/codegen_trn.cc | 13 +++--
src/backend/trn/codegen/codegen_trn.h | 1 +
src/backend/trn/op/target_builtin.cc | 7 +++
tests/python/tirx/codegen/test_codegen_nki.py | 28 +++++-----
.../operator/tile_primitive/trn/test_binary_trn.py | 12 ++--
.../tile_primitive/trn/test_compose_op_trn.py | 54 +++++++++---------
.../operator/tile_primitive/trn/test_copy_trn.py | 64 +++++++++++-----------
.../operator/tile_primitive/trn/test_gemm_trn.py | 30 +++++-----
.../tile_primitive/trn/test_private_alloc_trn.py | 10 ++--
.../tile_primitive/trn/test_reduction_trn.py | 18 +++---
.../operator/tile_primitive/trn/test_select_trn.py | 8 +--
.../operator/tile_primitive/trn/test_unary_trn.py | 14 ++---
24 files changed, 168 insertions(+), 147 deletions(-)
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 0e96a758df..6331344f52 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -882,7 +882,6 @@ constexpr const char* kVolatile = "tirx.volatile";
constexpr const char* buffer_data_alignment = "buffer_data_alignment";
/*! \brief Mark buffer allocated addr in bytes */
constexpr const char* buffer_allocated_addr = "buffer_allocated_addr";
-constexpr const char* tensorized_nki_instruction =
"tensorized_nki_instruction";
/*!
* \brief Mark the kernel as persistent.
diff --git a/python/tvm/backend/trn/script.py b/python/tvm/backend/trn/script.py
index f92dbd4b67..1bd1baee14 100644
--- a/python/tvm/backend/trn/script.py
+++ b/python/tvm/backend/trn/script.py
@@ -56,5 +56,14 @@ class NKINamespace:
self.identity = wrap(_trn_op.nki_identity)
self.affine_select = wrap(_trn_op.nki_affine_select)
+ @staticmethod
+ def tensorized_instruction():
+ """Tensorize the NKI instructions in the region body."""
+ from tvm.tirx.script.ir_builder import (
+ region, # pylint: disable=import-outside-toplevel
+ )
+
+ return region("tirx.nki.tensorized_instruction", [])
+
__all__ = ["NKINamespace", "OpWrapper"]
diff --git a/python/tvm/backend/trn/tile_primitive/binary/default.py
b/python/tvm/backend/trn/tile_primitive/binary/default.py
index 9cf669d892..7893a022ce 100644
--- a/python/tvm/backend/trn/tile_primitive/binary/default.py
+++ b/python/tvm/backend/trn/tile_primitive/binary/default.py
@@ -76,7 +76,7 @@ def binary_trn(
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, b_var: b_loop})
diff --git a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
index 7c29db3794..d2a5ffda5e 100644
--- a/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
@@ -95,7 +95,7 @@ def binary_chain_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, b_var: b_loop})
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 fd00bfc0a1..39fb66e417 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
@@ -105,7 +105,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, spatial_b_var: b_loop}) # noqa: E501
@@ -127,7 +127,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc
def impl():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size,
annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop,
f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop}) #
noqa: E501
@@ -139,7 +139,7 @@ def binary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc
T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop],
dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)],
src2[tuple(src_2_indices)], binary_opcode, reduce_opcode, reverse[0]) # noqa:
E501
else:
T.nki.tensorscalar_reduce(intermediate_buffer[p_loop, reduction_b_loop],
dst1[tuple(vec_dst_idx)], src1[tuple(src_1_indices)], CONST, binary_opcode,
reduce_opcode, reverse[0]) # noqa: E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, reduction_b_extent,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop,
spatial_b_var: b_loop})
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 8b56df56b1..ae8c720a3d 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
@@ -104,7 +104,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, spatial_b_var: b_loop}) # noqa: E501
@@ -131,7 +131,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
def impl():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size,
annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop,
f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop}) #
noqa: E501
@@ -143,7 +143,7 @@ def unary_reduce_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc |
T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop,
reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)],
unary_opcode, reduce_opcode, bias_buffer[tuple(src_bias_indices)], scale)) #
noqa: E501
else:
T.evaluate(T.nki.activation_reduce(intermediate_buffer[p_loop,
reduction_b_loop], dst1[tuple(dst_1_indices)], src[tuple(src_1_indices)],
unary_opcode, reduce_opcode, bias_buffer[p_loop, f_loop], scale)) # noqa: E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, reduction_b_extent,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop,
spatial_b_var: b_loop})
diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py
b/python/tvm/backend/trn/tile_primitive/copy/default.py
index 06dd2ebd89..a02907e1db 100644
--- a/python/tvm/backend/trn/tile_primitive/copy/default.py
+++ b/python/tvm/backend/trn/tile_primitive/copy/default.py
@@ -97,7 +97,7 @@ def transpose_schedule(
# This fragment captures buffers and indices from its insertion scope.
@T.prim_func(check_well_formed=False)
def identity_init():
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for rhs_f_loop in T.serial(0, rhs_f_size,
annotations={nki_dim: "F"}):
T.evaluate(T.nki.identity(identity_tensor[p_loop,
rhs_f_loop], p_size))
@@ -115,7 +115,7 @@ def transpose_schedule(
@T.prim_func(check_well_formed=False)
def transpose_psum_output():
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for lhs_f_loop in T.serial(0, lhs_f_size,
annotations={nki_dim: "lhs_F"}):
for rhs_f_loop in T.serial(
@@ -169,7 +169,7 @@ def transpose_schedule(
def transpose_sbuf_output():
for b_loop in T.serial(0, b_extent):
for extend_b_loop in T.serial(0, extend_len):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for lhs_f_loop in T.serial(0, lhs_f_size,
annotations={nki_dim: "lhs_F"}):
for rhs_f_loop in T.serial(0, rhs_f_size,
annotations={nki_dim: "rhs_F"}): # noqa: E501
@@ -177,7 +177,7 @@ def transpose_schedule(
src_indices =
T.meta_var(inst_gen.generate_indices(src_region))
if inst_gen.make_guard(src_region):
T.evaluate(T.nki.matmul(acc_psum[b_loop %
max_psum_slots, lhs_f_loop,extend_b_loop * rhs_f_size + rhs_f_loop],
src_buffer[tuple(src_indices)], identity_tensor[p_loop, rhs_f_loop])) # noqa:
E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, rhs_f_size * extend_len,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map(dst_region, {b_var: b_loop,
lhs_f: p_loop, dst_f: f_loop % rhs_f_size, extend_b: f_loop // rhs_f_size}) #
noqa: E501
@@ -275,7 +275,7 @@ def copy_trn(op: TilePrimitiveCall, sctx: DispatchContext)
-> PrimFunc | None:
def impl():
# the additional b loop is to satisfy hardware instuction size limit
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst.size, annotations={nki_dim:
"F"}):
inst_gen.set_bind_map_all({b_var: b_loop, p_var:
p_loop, f_var: f_loop})
diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py
b/python/tvm/backend/trn/tile_primitive/gemm/default.py
index 32f75ab4b5..73c439072c 100644
--- a/python/tvm/backend/trn/tile_primitive/gemm/default.py
+++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py
@@ -223,7 +223,7 @@ def matmul_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc | None:
# fmt: off
@T.inline
def matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop, acc,
C_as_output, max_psum_slots): # noqa: E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(0, lhs_f_size,
annotations={"nki_dim": "lhs_F"}):
for rhs_f_loop in T.serial(0, inst_repr.size,
annotations={"nki_dim": "rhs_F"}):
@@ -273,7 +273,7 @@ def matmul_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc | None:
for lhs_b_loop, rhs_b_loop in T.grid(lhs_b_extent, rhs_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
matmul_inst_macro(lhs_b_loop, rhs_b_loop, reduction_b_loop,
acc_psum, False, max_psum_slots) # noqa: E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for lhs_f_loop in T.serial(0, lhs_f_size,
annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(0, inst_repr.size,
annotations={"nki_dim": "F"}):
inst_gen.set_bind_map(C_buffer_region, {lhs_f:
lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop}) # noqa:
E501
diff --git a/python/tvm/backend/trn/tile_primitive/private_alloc.py
b/python/tvm/backend/trn/tile_primitive/private_alloc.py
index 7048afbd4c..ceafb52120 100644
--- a/python/tvm/backend/trn/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -70,7 +70,7 @@ def alloc_const_bias_trn(
# This fragment captures buffers and indices from its insertion scope.
@T.prim_func(check_well_formed=False)
def const_bias_init():
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, par_size, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, max_inst_size, annotations={nki_dim:
"F"}):
T.evaluate(T.nki.memset(new_buffer[p_loop, f_loop], bias))
@@ -121,7 +121,7 @@ def alloc_identity_trn(
# This fragment captures buffers and indices from its insertion scope.
@T.prim_func(check_well_formed=False)
def identity_init():
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, par_size, annotations={nki_dim: "P"}):
for rhs_f_loop in T.serial(0, par_size, annotations={nki_dim:
"F"}):
T.evaluate(T.nki.identity(new_buffer[p_loop, rhs_f_loop],
par_size))
diff --git a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
index b92ef4d8ca..3ca7277238 100644
--- a/python/tvm/backend/trn/tile_primitive/reduction/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
@@ -132,7 +132,7 @@ def reduction_trn(
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, spatial_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, spatial_b_var: b_loop}) # noqa: E501
@@ -148,14 +148,14 @@ def reduction_trn(
def two_stage_reduction():
for b_loop in T.serial(0, spatial_b_extent):
for reduction_b_loop in T.serial(0, reduction_b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size,
annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop,
f_var: f_loop, spatial_b_var: b_loop, reduction_b_var: reduction_b_loop}) #
noqa: E501
if inst_gen.make_guard(src_buffer_region):
src_indices =
T.meta_var(inst_gen.generate_indices(src_buffer_region)) # noqa: E501
T.evaluate(T.nki.tensorreduce(intermediate_buffer[p_loop, reduction_b_loop],
src[src_indices], opcode, False, -1)) # noqa: E501
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, reduction_b_extent,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map(src_buffer_region, {p_var:
p_loop, f_var: 0, spatial_b_var: b_loop, reduction_b_var: f_loop}) # noqa: E501
diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py
b/python/tvm/backend/trn/tile_primitive/select/default.py
index e8f40f4507..a65a937488 100644
--- a/python/tvm/backend/trn/tile_primitive/select/default.py
+++ b/python/tvm/backend/trn/tile_primitive/select/default.py
@@ -113,7 +113,7 @@ def select_trn(op: TilePrimitiveCall, sctx:
DispatchContext) -> PrimFunc | None:
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({f_var: f_loop, p_var:
p_loop, b_var: b_loop})
diff --git a/python/tvm/backend/trn/tile_primitive/unary/utils.py
b/python/tvm/backend/trn/tile_primitive/unary/utils.py
index 52df3cc0c2..bc86fc433e 100644
--- a/python/tvm/backend/trn/tile_primitive/unary/utils.py
+++ b/python/tvm/backend/trn/tile_primitive/unary/utils.py
@@ -108,7 +108,7 @@ def get_const_bias_tensor(bias, shape, dtype, workspace,
sctx):
# This fragment captures buffers and indices from its insertion scope.
@T.prim_func(check_well_formed=False)
def const_bias_init():
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, shape[0], annotations={nki_dim:
"P"}):
for f_loop in T.serial(0, shape[1], annotations={nki_dim:
"F"}):
T.evaluate(T.nki.memset(bias_buffer[p_loop, f_loop],
bias))
@@ -171,7 +171,7 @@ def generate_unary_func(
@T.prim_func(check_well_formed=False)
def impl():
for b_loop in T.serial(0, b_extent):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}):
for f_loop in T.serial(0, inst_repr.size,
annotations={nki_dim: "F"}):
inst_gen.set_bind_map_all({p_var: p_loop, f_var:
f_loop, b_var: b_loop})
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index 8b076a9598..6d611a0700 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -280,15 +280,20 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode*
op, const CallNode* bu
void CodeGenTrainium::Dispatch_(const AttrStmtNode* op) {
TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
<< "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- if (op->attr_key == tirx::attr::tensorized_nki_instruction) {
+ this->PrintStmt(op->body);
+}
+
+void CodeGenTrainium::Dispatch_(const RegionStmtNode* op) {
+ static const Op tensorized_instruction =
Op::Get("tirx.nki.tensorized_instruction");
+ if (op->op.same_as(tensorized_instruction)) {
ctx_.tensorizing = true;
ctx_.mask = std::nullopt;
ctx_.loopvar2dim.clear();
ctx_.is_matmul_input = false;
- }
- this->PrintStmt(op->body);
- if (op->attr_key == tirx::attr::tensorized_nki_instruction) {
+ this->PrintStmt(op->body);
ctx_.tensorizing = false;
+ } else {
+ CodeGenC::Dispatch_(op);
}
}
diff --git a/src/backend/trn/codegen/codegen_trn.h
b/src/backend/trn/codegen/codegen_trn.h
index c7b1d21928..5f1d16441f 100644
--- a/src/backend/trn/codegen/codegen_trn.h
+++ b/src/backend/trn/codegen/codegen_trn.h
@@ -61,6 +61,7 @@ class CodeGenTrainium final : public CodeGenC {
void Dispatch_(const BindNode* op) final;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const AttrStmtNode* op) final; //
NOLINT(*)
+ void Dispatch_(const RegionStmtNode* op) final; //
NOLINT(*)
void Dispatch_(const ForNode* op) final; //
NOLINT(*)
void Dispatch_(const BufferStoreNode* op) final; //
NOLINT(*)=
void Dispatch_(const EvaluateNode* op) final; //
NOLINT(*)
diff --git a/src/backend/trn/op/target_builtin.cc
b/src/backend/trn/op/target_builtin.cc
index 7d9e835e92..439e47dab6 100644
--- a/src/backend/trn/op/target_builtin.cc
+++ b/src/backend/trn/op/target_builtin.cc
@@ -47,6 +47,13 @@ void RegisterTRNTargetBuiltins() {
registered = true;
RegisterNKIIntrinsicAliases();
+
+ OpDef("tirx.nki.tensorized_instruction", "Tensorize the NKI instructions in
the region body.")
+ .signature()
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("device_intrin"))
+ .set_attr<TDeviceIntrinsicNamespace>("TDeviceIntrinsicNamespace",
ffi::String("nki"))
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
+
ffi::String("tirx.nki.tensorized_instruction"));
}
namespace {
diff --git a/tests/python/tirx/codegen/test_codegen_nki.py
b/tests/python/tirx/codegen/test_codegen_nki.py
index 043cd73d2f..fed427db5d 100644
--- a/tests/python/tirx/codegen/test_codegen_nki.py
+++ b/tests/python/tirx/codegen/test_codegen_nki.py
@@ -44,15 +44,15 @@ def test_nki_add_1():
T.device_entry()
A_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
B_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.load(A_sbuf[i, j], A[i, j])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.tensorscalar(B_sbuf[i, j], A_sbuf[i, j],
T.float32(1.0), "add")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.store(B[i, j], B_sbuf[i, j])
@@ -99,15 +99,15 @@ def test_nki_add_2():
A_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
B_sbuf = T.alloc_tensor((128, 512), "float32", scope="trn.sbuf",)
for k in range(0, 4):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.load(A_sbuf[i, j], A[i, 512*k+j])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.tensorscalar(B_sbuf[i, j], A_sbuf[i, j],
T.float32(1.0), "add")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(0, 128):
for j in range(0, 512):
T.nki.store(B[i, 512*k+j], B_sbuf[i, j])
@@ -187,7 +187,7 @@ def test_nki_matmul_1():
res_tile = T.alloc_tensor((1, TILE_M, TILE_N), "float32",
scope="trn.psum")
result_packed = T.alloc_tensor((TILE_K, BLOCK_N), "float32",
scope="trn.sbuf")
for n in range(NUM_BLOCK_N):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i0 in range(TILE_M):
for i1 in range(NUM_BLOCK_M):
for i2 in range(TILES_IN_BLOCK_M):
@@ -196,7 +196,7 @@ def test_nki_matmul_1():
T.nki.memset(result_tiles[i0, i1, i2, i3,
i4], T.float32(0.0))
for k in range(NUM_BLOCK_K):
for bk_r in range(TILES_IN_BLOCK_K):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_K):
for j in range(BLOCK_N):
T.nki.load(
@@ -208,7 +208,7 @@ def test_nki_matmul_1():
)
for m in range(NUM_BLOCK_M):
for bk_l in range(TILES_IN_BLOCK_K):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_K):
for j in range(BLOCK_M):
T.nki.load(
@@ -220,12 +220,12 @@ def test_nki_matmul_1():
)
for bn in range(TILES_IN_BLOCK_N):
for bm in range(TILES_IN_BLOCK_M):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_M):
for j in range(TILE_N):
T.nki.memset(res_tile[0, i, j],
T.float32(0.0))
for bk in range(TILES_IN_BLOCK_K):
- with T.attr(0, "tensorized_nki_instruction",
1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_M):
for j in range(TILE_N):
for k in range(TILE_K):
@@ -235,7 +235,7 @@ def test_nki_matmul_1():
rhs_tiles[k, bk, bn *
TILE_N + j],
1,
)
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_M):
for j in range(TILE_N):
T.nki.tensortensor(
@@ -247,14 +247,14 @@ def test_nki_matmul_1():
for m in range(NUM_BLOCK_M):
for bm in range(TILES_IN_BLOCK_M):
for bn in range(TILES_IN_BLOCK_N):
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_K):
for j in range(TILE_N):
T.nki.tensor_copy(
result_packed[i, bn * TILE_N + j],
result_tiles[i, m, bm, bn, j],
)
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for i in range(TILE_K):
for j in range(BLOCK_N):
T.nki.store(
diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.py
index 5a15976772..63028e7da7 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_binary_trn.py
@@ -95,7 +95,7 @@ def test_simple_binary(op_type, operands_type):
B_sbuf = T.alloc_tensor(src2_shape, scope="trn.sbuf")
C_sbuf = T.alloc_tensor(dst_shape, scope="trn.sbuf")
for b_loop in T.serial(0, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if T.constexpr(operands_type == "region_region"):
@@ -182,7 +182,7 @@ def test_binary_complex(op_type, operands_type):
B_sbuf_view = T.decl_tensor(src2_layout_data_iter, data=B_sbuf.data,
scope="trn.sbuf", layout=None) # noqa: E501
C_sbuf_view = T.decl_tensor((128, 2048), data=C_sbuf.data,
scope="trn.sbuf", layout=None)
for i, b_loop in T.grid(4, b_extent):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, f_extent,
annotations={"nki_dim":"F"}):
if T.constexpr(operands_type == "region_region"):
@@ -228,7 +228,7 @@ def test_binary_broadcast1():
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for b_loop in T.serial(0, 512):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorscalar(C_sbuf[p_loop, b_loop % 4 * 4096 +
b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 * 32
+ f_loop], B_sbuf[p_loop, b_loop], "add", T.bool(False)) # noqa: E501
@@ -264,7 +264,7 @@ def test_binary_broadcast2():
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for b_loop in T.serial(0, 128):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
T.nki.tensortensor(C_sbuf[p_loop, b_loop % 4 * 4096 +
b_loop // 4 * 128 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop // 4 *
128 + f_loop], B_sbuf[p_loop, b_loop % 4 * 128 + f_loop], "add") # noqa: E501
@@ -300,7 +300,7 @@ def test_binary_broadcast3():
B_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in T.serial(0, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
T.nki.tensortensor(C_sbuf[p_loop, b_loop * 128 + f_loop],
A_sbuf[p_loop, b_loop * 128 + f_loop], B_sbuf[p_loop, b_loop * 4096 + f_loop],
"add") # noqa: E501
@@ -337,7 +337,7 @@ def test_binary_with_guard():
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for j, b_loop in T.grid(4, 96):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
if b_loop % 3 - j < 0:
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 881fd8735b..c292f06169 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
@@ -70,7 +70,7 @@ def test_simple_activation_reduce():
def expected():
T.func_attr({"global_symbol": "activation_reduce"})
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -78,7 +78,7 @@ def test_simple_activation_reduce():
B = T.alloc_tensor((128, 512), scope="trn.sbuf")
C = T.alloc_tensor((128, 1), scope="trn.sbuf")
for b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.activation_reduce(C[p_loop, 0], B[p_loop, f_loop],
A[p_loop, f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop]) # noqa: E501
@@ -112,7 +112,7 @@ def test_activation_reduce_in_loop():
def expected():
T.func_attr({"global_symbol": "activation_reduce"})
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -120,7 +120,7 @@ def test_activation_reduce_in_loop():
B = T.alloc_tensor((128, 8192), scope="trn.sbuf")
C = T.alloc_tensor((128, 16), scope="trn.sbuf")
for i, b_loop in T.grid(2, 16):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 +
b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop % 8 // 2 * 2048 + b_loop // 8 *
1024 + b_loop % 2 * 512 + f_loop], A[p_loop, i * 8192 + b_loop * 512 + f_loop],
"sqrt", "add", bias=const_bias[p_loop, f_loop]) # noqa: E501
@@ -154,7 +154,7 @@ def test_activation_reduce_in_loop2():
def expected():
T.func_attr({"global_symbol": "activation_reduce"})
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -162,7 +162,7 @@ def test_activation_reduce_in_loop2():
B = T.alloc_tensor((128, 8192), scope="trn.sbuf")
C = T.alloc_tensor((128, 16), scope="trn.sbuf")
for i, b_loop in T.grid(2, 16):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 +
b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i *
8192 + b_loop * 512 + f_loop], "sqrt", "add", bias=const_bias[p_loop, f_loop])
# noqa: E501
@@ -197,7 +197,7 @@ def test_activation_reduce_two_stage():
T.func_attr({"global_symbol": "activation_reduce"})
partial_reduce = T.alloc_tensor((128, 8), scope="trn.sbuf")
const_bias = T.alloc_tensor((128, 1024), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -206,11 +206,11 @@ def test_activation_reduce_two_stage():
C = T.alloc_tensor((128, 1), scope="trn.sbuf")
for i, b_loop in T.grid(2, 1):
for reduction_b_loop in range(8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.activation_reduce(partial_reduce[p_loop,
reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop //
4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop],
"sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
T.nki.tensorreduce(C[p_loop, 0], partial_reduce[p_loop,
f_loop], "add", T.bool(False), -1) # noqa: E501
@@ -251,7 +251,7 @@ def test_activation_reduce_with_bias_scale():
C = T.alloc_tensor((128, 16), scope="trn.sbuf")
bias = T.alloc_tensor((128, 1), scope="trn.sbuf")
for i, b_loop in T.grid(2, 16):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.activation_reduce(C[p_loop, b_loop % 8 // 2 * 4 +
b_loop // 8 * 2 + b_loop % 2], B[p_loop, b_loop * 512 + f_loop], A[p_loop, i *
8192 + b_loop * 512 + f_loop], "sqrt", "add", bias[p_loop, 0], T.float32(2.0))
# noqa: E501
@@ -286,7 +286,7 @@ def test_simple_tensor_scalar_reduce():
B = T.alloc_tensor((128, 512), scope="trn.sbuf")
C = T.alloc_tensor((128, 1), scope="trn.sbuf")
for b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.tensorscalar_reduce(C[p_loop, 0], B[p_loop, f_loop],
A[p_loop, f_loop], T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
@@ -352,7 +352,7 @@ def test_tensor_scalar_reduce_complex():
C_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
D_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in range(512):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorscalar_reduce(D_sbuf[p_loop, b_loop % 4 * 128
+ b_loop // 4], C_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4],
A_sbuf[p_loop, b_loop % 4 * 4096 + f_loop * 128 + b_loop // 4], B_sbuf[p_loop,
b_loop % 4 * 128 + b_loop // 4], "add", "add", T.bool(True)) # noqa: E501
@@ -389,11 +389,11 @@ def test_tensor_scalar_reduce_two_stage():
C_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for b_loop in range(4):
for reduction_b_loop in range(4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.tensorscalar_reduce(partial_reduce[p_loop,
reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 +
f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop],
T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
T.nki.tensorreduce(C_sbuf[p_loop, b_loop],
partial_reduce[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
@@ -435,7 +435,7 @@ def test_vector_chain():
D_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
E_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for b_loop in T.serial(0, 512):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.scalar_tensor_scalar(E_sbuf[p_loop, b_loop % 4 *
4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop //
4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4], "add",
"add", T.bool(False), T.bool(True)) # noqa: E501
@@ -477,7 +477,7 @@ def test_vector_chain_2():
D_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
E_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for b_loop in T.serial(0, 512):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.scalar_tensor_tensor(E_sbuf[p_loop, b_loop % 4 *
4096 + b_loop // 4 * 32 + f_loop], A_sbuf[p_loop, b_loop % 4 * 4096 + b_loop //
4 * 32 + f_loop], B_sbuf[p_loop, b_loop], D_sbuf[p_loop, b_loop % 4 * 4096 +
b_loop // 4 * 32 + f_loop], "add", "add", T.bool(False), T.bool(True)) # noqa:
E501
@@ -510,7 +510,7 @@ def test_reduce_negate():
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop,
f_loop * 4 + i], "add", True, -1) # noqa: E501
@@ -547,7 +547,7 @@ def test_binary_reduce_guard():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for j, i, b_loop in T.grid(4, 4, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if b_loop - j < 1 and f_loop < i * 128 + 128:
@@ -583,7 +583,7 @@ def test_unary_reduce_guard():
def expected():
T.func_attr({"global_symbol": "unary_reduce"})
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -591,7 +591,7 @@ def test_unary_reduce_guard():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for j, i, b_loop in T.grid(4, 4, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
if b_loop - j < 1 and f_loop < i * 128 + 128:
@@ -631,7 +631,7 @@ def test_binary_chain_guard():
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for j, i, b_loop in T.grid(4, 4, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if b_loop - j < 1 and f_loop < i * 128 + 128:
@@ -667,7 +667,7 @@ def test_activation_reduce_two_stage_workspace():
def expected():
T.func_attr({"global_symbol": "activation_reduce"})
const_bias = T.alloc_tensor((128, 1024), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -677,11 +677,11 @@ def test_activation_reduce_two_stage_workspace():
C = T.alloc_tensor((128, 1), scope="trn.sbuf")
for i, b_loop in T.grid(2, 1):
for reduction_b_loop in range(8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.activation_reduce(intermediate_buffer[p_loop,
reduction_b_loop], B[p_loop, reduction_b_loop % 4 * 2048 + reduction_b_loop //
4 * 1024 + f_loop], A[p_loop, i * 8192 + reduction_b_loop * 1024 + f_loop],
"sqrt", "add", const_bias[p_loop, f_loop], T.float32(1.0)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(8, annotations={"nki_dim": "F"}):
T.nki.tensorreduce(C[p_loop, 0],
intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
@@ -721,11 +721,11 @@ def test_tensor_scalar_reduce_two_stage_workspace():
C_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for b_loop in range(4):
for reduction_b_loop in range(4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.tensorscalar_reduce(intermediate_buffer[p_loop,
reduction_b_loop], B_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 +
f_loop], A_sbuf[p_loop, reduction_b_loop * 4096 + b_loop * 1024 + f_loop],
T.float32(1.0), "add", "add", T.bool(False)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(4, annotations={"nki_dim": "F"}):
T.nki.tensorreduce(C_sbuf[p_loop, b_loop],
intermediate_buffer[p_loop, f_loop], "add", T.bool(False), -1) # noqa: E501
@@ -756,7 +756,7 @@ def test_unary_reduce_complex():
qk = T.alloc_tensor((128, 16384), scope="trn.sbuf")
running_max = T.alloc_tensor((128, 128), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(8192, annotations={"nki_dim": "F"}):
T.nki.activation_reduce(rowsum_p[p_loop, i % 2], p[p_loop,
f_loop], qk[p_loop, i % 2 * 8192 + f_loop], "exp", "add", running_max[p_loop,
i], T.float32(1.0)) # noqa: E501
diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py
index c34fb6d5de..f279a20e75 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_copy_trn.py
@@ -67,7 +67,7 @@ def test_simple_copy():
A_1 = T.decl_tensor((65536,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in T.serial(0, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.load(A_sbuf[p_loop, f_loop], A_1[p_loop * 512 +
f_loop])
@@ -98,7 +98,7 @@ def test_simple_copy_2():
A_1 = T.decl_tensor((65536,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in T.serial(0, 512):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 1, annotations={"nki_dim": "F"}):
T.nki.load(A_sbuf[p_loop, b_loop], A_1[b_loop * 128 +
p_loop])
@@ -129,7 +129,7 @@ def test_copy_in_a_loop():
A_1 = T.decl_tensor((262144,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.load(
@@ -166,7 +166,7 @@ def test_copy_in_a_loop_2():
A_sbuf_view = T.decl_tensor((128, 2048), data=A_sbuf.data,
scope="trn.sbuf", layout=None)
A_view = T.decl_tensor((262144,), data=_A_flat.data, layout=None)
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.load(
@@ -200,7 +200,7 @@ def test_copy_transpose():
T.func_attr({"global_symbol": "copy"})
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((8, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -208,12 +208,12 @@ def test_copy_transpose():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for b_loop in range(16):
for extend_b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(128, annotations={"nki_dim":
"lhs_F"}):
for rhs_f_loop in T.serial(128,
annotations={"nki_dim": "rhs_F"}):
T.nki.matmul(acc_psum[b_loop % 8, lhs_f_loop,
rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop], identity[p_loop,
rhs_f_loop], T.bool(True)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop],
acc_psum[b_loop % 8, p_loop, f_loop]) # noqa: E501
@@ -247,7 +247,7 @@ def test_copy_transpose_2():
T.func_attr({"global_symbol": "copy"})
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((8, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -256,12 +256,12 @@ def test_copy_transpose_2():
for i in range(4):
for b_loop in range(4):
for extend_b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(128,
annotations={"nki_dim": "lhs_F"}):
for rhs_f_loop in T.serial(128,
annotations={"nki_dim": "rhs_F"}):
T.nki.matmul(acc_psum[b_loop, lhs_f_loop,
rhs_f_loop], A_sbuf[p_loop, lhs_f_loop * 4 + b_loop], identity[p_loop,
rhs_f_loop], T.bool(True)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + i * 4 +
b_loop], acc_psum[b_loop, p_loop, f_loop]) # noqa: E501
@@ -293,7 +293,7 @@ def test_copy_different_f():
A_sbuf = T.alloc_tensor((128, 256), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 256), scope="trn.sbuf")
for b_loop in T.serial(0, 64):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(
@@ -331,7 +331,7 @@ def test_copy_different_shape():
B_sbuf = T.alloc_tensor((128, 16), scope="trn.sbuf")
B_sbuf_view = T.decl_tensor((128, 16), data=B_sbuf.data,
scope="trn.sbuf", layout=None)
for b_loop in T.serial(0, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 4, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(
@@ -365,7 +365,7 @@ def test_copy_irregular_shape():
A_1 = T.decl_tensor((1280000,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.store(A_1[p_loop * 10000 + i * 512 + f_loop],
A_sbuf[p_loop, f_loop])
@@ -398,7 +398,7 @@ def test_copy_different_shape_dim():
A_1 = T.decl_tensor((2097152,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for i, b_loop in T.grid(32, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.load(A_sbuf[p_loop, f_loop], A_1[i * 65536 + p_loop
* 128 + f_loop])
@@ -429,7 +429,7 @@ def test_copy_with_offset():
A_1 = T.decl_tensor((131072,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for i, b_loop in T.grid(2, 2):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.load(
@@ -463,7 +463,7 @@ def test_large_dma_copy():
A_1 = T.decl_tensor((2097152,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 4096, annotations={"nki_dim": "F"}):
T.nki.load(
@@ -497,7 +497,7 @@ def test_copy_with_inst_size_limit():
B_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
A_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
for i, b_loop in T.grid(4, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(
@@ -532,7 +532,7 @@ def test_copy_with_complex_index():
A_1 = T.decl_tensor((16777216,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 32768), scope="trn.sbuf")
for b_loop in T.serial(0, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
T.nki.load(A_sbuf[p_loop, b_loop * 2048 + f_loop + 16384],
A_1[b_loop * 524288 + p_loop * 4096 + f_loop + 12584960]) # noqa: E501
@@ -564,7 +564,7 @@ def test_copy_with_complex_index_2():
A_1 = T.decl_tensor((4194304,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 131072), scope="trn.sbuf")
for b_loop in T.serial(0, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
T.nki.load(A_sbuf[p_loop, b_loop * 4096 + f_loop +
100352], A_1[b_loop * 262144 + p_loop * 2048 + f_loop + 2097152]) # noqa: E501
@@ -590,7 +590,7 @@ def test_copy_transpose_with_workspace():
B_sbuf = T.alloc_tensor(dst_shape, "float32", scope="trn.sbuf",
layout=dst_layout)
identity = T.alloc_tensor((128, 128), "float32", scope="trn.sbuf")
acc_psum = T.alloc_tensor((1, 128, 512), "float32", scope="trn.psum",
allocated_addr=(0, 0))
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -603,18 +603,18 @@ def test_copy_transpose_with_workspace():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((1, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
for b_loop in range(16):
for extend_b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(128, annotations={"nki_dim":
"lhs_F"}):
for rhs_f_loop in T.serial(128,
annotations={"nki_dim": "rhs_F"}):
T.nki.matmul(acc_psum[0, lhs_f_loop, extend_b_loop
* 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 128 + lhs_f_loop],
identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(B_sbuf[p_loop, f_loop * 16 + b_loop],
acc_psum[0, p_loop, f_loop]) # noqa: E501
@@ -648,7 +648,7 @@ def test_copy_with_guard():
A_1 = T.decl_tensor((262144,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for j, i, b_loop in T.grid(4, 4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
if f_loop < j * 128:
@@ -684,7 +684,7 @@ def test_copy_with_guard_2():
A_1 = T.decl_tensor((262144,), data=A.data, layout=None)
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for j, i, b_loop in T.grid(4, 4, 3):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 384, annotations={"nki_dim":"F"}):
if b_loop - j < 0 and f_loop < i * 128:
@@ -718,7 +718,7 @@ def test_copy_transpose_with_guard():
T.func_attr({"global_symbol": "copy"})
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((8, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -726,13 +726,13 @@ def test_copy_transpose_with_guard():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for i, j, b_loop in T.grid(4, 4, 3):
for extend_b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(128, annotations={"nki_dim":
"lhs_F"}):
for rhs_f_loop in T.serial(128,
annotations={"nki_dim": "rhs_F"}):
if b_loop - j < 0:
T.nki.matmul(acc_psum[b_loop, lhs_f_loop,
rhs_f_loop], A_sbuf[p_loop, i * 512 + b_loop * 128 + lhs_f_loop],
identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
if b_loop - j < 0:
@@ -766,7 +766,7 @@ def test_copy_with_specified_max_inst_size():
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf", layout=None)
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf", layout=None)
for b_loop in T.serial(0, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(A_sbuf[p_loop, b_loop * 128 + f_loop],
B_sbuf[p_loop, b_loop * 128 + f_loop]) # noqa: E501
@@ -791,7 +791,7 @@ def test_copy_transpose_with_extended_f():
T.func_attr({"global_symbol": "copy"})
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((8, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -799,12 +799,12 @@ def test_copy_transpose_with_extended_f():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for b_loop in range(4):
for extend_b_loop in range(4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for lhs_f_loop in T.serial(128, annotations={"nki_dim":
"lhs_F"}):
for rhs_f_loop in T.serial(128,
annotations={"nki_dim": "rhs_F"}):
T.nki.matmul(acc_psum[b_loop, lhs_f_loop,
extend_b_loop * 128 + rhs_f_loop], A_sbuf[p_loop, b_loop * 512 + extend_b_loop
* 128 + lhs_f_loop], identity[p_loop, rhs_f_loop], T.bool(True)) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.tensor_copy(B_sbuf[p_loop, b_loop * 512 + f_loop],
acc_psum[b_loop, p_loop, f_loop]) # noqa: E501
diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py
index eae318588b..4199e0d312 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_gemm_trn.py
@@ -70,7 +70,7 @@ def test_simple_gemm():
B_sbuf = T.alloc_tensor((128, 128), scope="trn.sbuf")
C_psum = T.alloc_tensor((1, 128, 128), scope="trn.psum")
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(1, 1, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
@@ -104,7 +104,7 @@ def test_larger_gemm():
B_sbuf = T.alloc_tensor((128, 1024), scope="trn.sbuf")
C_psum = T.alloc_tensor((1, 128, 512), scope="trn.psum")
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 1, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 256,
annotations={"nki_dim":"rhs_F"}):
@@ -145,7 +145,7 @@ def test_gemm_in_a_loop():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2,
1, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 256,
annotations={"nki_dim":"rhs_F"}):
@@ -186,7 +186,7 @@ def test_gemm_with_stride():
B_sbuf = T.alloc_tensor((128, 4095), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2,
1, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 256,
annotations={"nki_dim":"rhs_F"}):
@@ -228,7 +228,7 @@ def test_gemm_swap_lhs_rhs():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2,
2, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
@@ -270,12 +270,12 @@ def test_gemm_with_sbuf_output():
C_sbuf = T.alloc_tensor((128, 1024), scope="trn.sbuf")
for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
for reduction_b_loop in range(4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
T.nki.matmul(buffer[lhs_b_loop * 2 + rhs_b_loop,
lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 +
lhs_b_loop * 128 + lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k
* 512 + reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop *
256 + lhs_b_loop * 128 + rhs_f_loop], buffer[lhs_b_loop * 2 + rhs_b_loop,
lhs_f_loop, rhs_f_loop]) # noqa: E501
@@ -317,7 +317,7 @@ def test_gemm_different_shape():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2,
2, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
@@ -351,7 +351,7 @@ def test_gemm_too_large_f_size():
B_sbuf = T.alloc_tensor((128, 1024), scope="trn.sbuf")
C_psum = T.alloc_tensor((4, 128, 512), scope="trn.psum")
for lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 512,
annotations={"nki_dim":"rhs_F"}):
@@ -395,12 +395,12 @@ def test_gemm_sbuf_output_with_workspace():
C_psum = T.alloc_tensor((1, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
for i, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2):
for reduction_b_loop in range(4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
T.nki.matmul(C_psum[0, lhs_f_loop, rhs_f_loop],
B_sbuf[p_loop, k * 1024 + reduction_b_loop * 256 + lhs_b_loop * 128 +
lhs_f_loop], A_sbuf[p_loop, i * 2048 + rhs_b_loop * 1024 + k * 512 +
reduction_b_loop * 128 + rhs_f_loop], True) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
T.nki.tensor_copy(C_sbuf[lhs_f_loop, i * 512 + rhs_b_loop *
256 + lhs_b_loop * 128 + rhs_f_loop], C_psum[0, lhs_f_loop, rhs_f_loop]) #
noqa: E501
@@ -471,7 +471,7 @@ def test_gemm_transpose_AB():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(2, 2, 2,
1, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 256,
annotations={"nki_dim":"rhs_F"}):
@@ -515,13 +515,13 @@ def test_gemm_guard():
C_sbuf = T.alloc_tensor((128, 1024), scope="trn.sbuf")
for i, j, k, lhs_b_loop, rhs_b_loop in T.grid(2, 2, 2, 2, 2):
for reduction_b_loop in range(8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"rhs_F"}):
if reduction_b_loop - k * 4 < 4 and lhs_b_loop - j < 1
and 0 < i and reduction_b_loop - k * 4 < 4: # noqa: E501
T.nki.matmul(acc_psum[lhs_b_loop * 2 + rhs_b_loop,
lhs_f_loop, rhs_f_loop], B_sbuf[p_loop, reduction_b_loop * 256 + lhs_b_loop *
128 + lhs_f_loop], A_sbuf[p_loop, rhs_b_loop * 1024 + reduction_b_loop * 128 +
rhs_f_loop], True) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for lhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for rhs_f_loop in T.serial(0, 128, annotations={"nki_dim":"F"}):
if 0 < i and lhs_b_loop - j < 1:
@@ -564,7 +564,7 @@ def test_gemm_guard2():
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
C_psum = T.alloc_tensor((2, 128, 512), scope="trn.psum")
for j, i, k, lhs_b_loop, rhs_b_loop, reduction_b_loop in T.grid(4, 2,
2, 2, 1, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for lhs_f_loop in T.serial(0, 128,
annotations={"nki_dim":"lhs_F"}):
for rhs_f_loop in T.serial(0, 256,
annotations={"nki_dim":"rhs_F"}):
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 cf8f86145d..21082e0a6e 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
@@ -46,7 +46,7 @@ def test_copy_transpose():
T.device_entry()
identity = T.alloc_tensor((128, 128), scope="trn.sbuf")
acc_psum = T.alloc_tensor((8, 128, 512), scope="trn.psum",
allocated_addr=[0, 0])
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for rhs_f_loop in T.serial(128, annotations={"nki_dim": "F"}):
T.nki.identity(identity[p_loop, rhs_f_loop], 128)
@@ -104,7 +104,7 @@ def test_unary_with_bias_scale():
T.func_attr({"global_symbol": "unary"})
T.device_entry()
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
@@ -253,7 +253,7 @@ def test_activation_reduce_two_stage():
T.device_entry()
partial_reduce = T.alloc_tensor((128, 8), scope="trn.sbuf")
const_bias = T.alloc_tensor((128, 1024), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -296,7 +296,7 @@ def test_partial_workspace_specify():
T.func_attr({"global_symbol": "activation_reduce"})
T.device_entry()
const_bias = T.alloc_tensor((128, 1024), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
@@ -337,7 +337,7 @@ def test_workspace_reuse():
T.func_attr({"global_symbol": "unary"})
T.device_entry()
const_bias = T.alloc_tensor((128, 1024), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(1024, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(0.0))
diff --git
a/tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py
index dab2091f70..28a30c663d 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_reduction_trn.py
@@ -77,7 +77,7 @@ def test_simple_reduction(op_type):
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 1), scope="trn.sbuf")
for b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop,
f_loop], opcode, False, -1)
@@ -109,7 +109,7 @@ def test_reduction_with_multiple_axes():
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 1), scope="trn.sbuf")
for b_loop in range(1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 2048, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, 0], A_sbuf[p_loop,
f_loop], "add", False, -1)
@@ -142,7 +142,7 @@ def test_reduction_in_loop():
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, i], A_sbuf[p_loop,
f_loop * 4 + i], "add", False, -1) # noqa: E501
@@ -175,11 +175,11 @@ def test_reduction_two_stage():
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for b_loop in range(4):
for reduction_b_loop in range(32):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(intermediate_buffer[p_loop,
reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 +
f_loop], "add", False, -1) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, b_loop],
intermediate_buffer[p_loop, f_loop], "add", False, -1) # noqa: E501
@@ -217,7 +217,7 @@ def test_reduction_with_guard():
for i, j in T.grid(4, 4):
for b_loop in range(4):
for reduction_b_loop in range(2):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim":
"F"}):
if (
@@ -225,7 +225,7 @@ def test_reduction_with_guard():
and reduction_b_loop * 512 + f_loop < j * 256
+ 256
):
T.nki.tensorreduce(intermediate_buffer[p_loop,
reduction_b_loop], A_sbuf[p_loop, b_loop * 2048 + reduction_b_loop * 512 +
f_loop], "add", T.bool(False), -1) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(2, annotations={"nki_dim": "F"}):
if b_loop - i < 1 and f_loop * 2 - j < 1:
@@ -262,11 +262,11 @@ def test_reduction_two_stage_workspace():
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
for b_loop in range(4):
for reduction_b_loop in range(32):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(intermediate_buffer[p_loop,
reduction_b_loop], A_sbuf[p_loop, reduction_b_loop * 128 + b_loop * 32 +
f_loop], "add", False, -1) # noqa: E501
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 32, annotations={"nki_dim":"F"}):
T.nki.tensorreduce(B_sbuf[p_loop, b_loop],
intermediate_buffer[p_loop, f_loop], "add", False, -1) # noqa: E501
diff --git a/tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py
b/tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py
index bf5cf5527a..341598eb6b 100644
--- a/tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py
+++ b/tests/python/tirx/operator/tile_primitive/trn/test_select_trn.py
@@ -68,7 +68,7 @@ def test_select():
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in T.serial(0, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.affine_select(B_sbuf[p_loop, f_loop], p_loop <
f_loop, A_sbuf[p_loop, f_loop], T.float32(0.0)) # noqa: E501
@@ -102,7 +102,7 @@ def test_select_in_loop():
A_sbuf = T.alloc_tensor((128, 16384), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for i, b_loop in T.grid(2, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.affine_select(B_sbuf[p_loop, f_loop], (i + 1) *
p_loop < f_loop, A_sbuf[p_loop, i * 8192 + f_loop], T.float32(0.0)) # noqa:
E501
@@ -135,7 +135,7 @@ def test_select_expr_affine():
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for b_loop in T.serial(0, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.affine_select(B_sbuf[p_loop, b_loop * 512 + f_loop],
b_loop * 128 + p_loop < f_loop, A_sbuf[p_loop, b_loop * 512 + f_loop],
T.float32(0.0)) # noqa: E501
@@ -169,7 +169,7 @@ def test_select_with_guard():
A_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 2048), scope="trn.sbuf")
for i, j, b_loop in T.grid(4, 4, 4):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if b_loop - i < 1 and f_loop < j * 128 + 128:
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 8226d9d2b0..e3287afd3d 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
@@ -76,7 +76,7 @@ def test_simple_unary(op_type):
A_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
B_sbuf = T.alloc_tensor((128, 512), scope="trn.sbuf")
for b_loop in T.serial(0, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if T.constexpr(op_type == "reciprocal"):
@@ -123,7 +123,7 @@ def test_unary_in_a_loop(op_type):
A_sbuf_view = T.decl_tensor((128, 4096), data=A_sbuf.data,
scope="trn.sbuf", layout=None)
B_sbuf_view = T.decl_tensor((128, 2048), data=B_sbuf.data,
scope="trn.sbuf", layout=None)
for i, b_loop in T.grid(4, 1):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if T.constexpr(op_type == "reciprocal"):
@@ -153,7 +153,7 @@ def test_unary_complex1():
T.func_attr({"global_symbol": "unary"})
A_sbuf = T.alloc_tensor((128, 8192), scope="trn.sbuf")
for b_loop in T.serial(0, 16):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.memset(A_sbuf[p_loop, b_loop * 512 + f_loop],
T.float32(0.0))
@@ -191,7 +191,7 @@ def test_unary_with_bias_scale(op_type):
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 4096), scope="trn.sbuf")
for b_loop in T.serial(0, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop],
A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, B_sbuf[p_loop, b_loop//2],
T.float32(2.0)) # noqa: E501
@@ -224,14 +224,14 @@ def test_unary_with_bias_scale_2(op_type):
def expected():
T.func_attr({"global_symbol": "unary"})
const_bias = T.alloc_tensor((128, 512), scope="trn.sbuf")
- with T.attr(0, "tensorized_nki_instruction", 1):
+ with T.nki.tensorized_instruction():
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.memset(const_bias[p_loop, f_loop], T.float32(1.0))
A_sbuf = T.alloc_tensor((128, 4096), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 4096), scope="trn.sbuf")
for b_loop in T.serial(0, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(128, annotations={"nki_dim": "P"}):
for f_loop in T.serial(512, annotations={"nki_dim": "F"}):
T.nki.activation(C_sbuf[p_loop, b_loop * 512 + f_loop],
A_sbuf[p_loop, b_loop * 512 + f_loop], op_type, const_bias[p_loop, f_loop],
T.float32(2.0)) # noqa: E501
@@ -270,7 +270,7 @@ def test_unary_with_guard():
B_sbuf = T.alloc_tensor((128, 4), scope="trn.sbuf")
C_sbuf = T.alloc_tensor((128, 4096), scope="trn.sbuf")
for i, j, b_loop in T.grid(4, 4, 8):
- T.attr(0, "tensorized_nki_instruction", 1)
+ T.nki.tensorized_instruction()
for p_loop in T.serial(0, 128, annotations={"nki_dim":"P"}):
for f_loop in T.serial(0, 512, annotations={"nki_dim":"F"}):
if b_loop // 2 - i < 1 and b_loop % 2 * 512 + f_loop < j *
256 + 256:
