This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 cb80add495 fix(lower-tirx): restore CUDA lowering with native IR
builders (#20620)
cb80add495 is described below
commit cb80add495bb4ebefee0d918c14c9408f1704ea2
Author: Bohan Hou <[email protected]>
AuthorDate: Thu Oct 8 18:26:11 2026 -0400
fix(lower-tirx): restore CUDA lowering with native IR builders (#20620)
The native TIRx builder migration leaves CUDA tile primitives using
implicit host branches, TensorMap calls without a native attribute
schema, and TMEM types whose address metadata references replaced
variables. This change restores lowering against the native IR
contracts.
- Mark compile-time control flow explicitly and keep constant TMEM
coordinates visible to address folding.
- Register TensorMap attributes and defaults, and remap `allocated_addr`
when rebuilding tensor types.
- Allow dispatch fragments to capture enclosing thread bindings while
retaining validation of complete functions; reject distributed tensors
in scalar copy fallback and normalize single-axis DSMEM loops.
- Add regressions for TensorMap attributes, TMEM metadata, and fallback
rejection. Update instruction/IR assertions and remove three obsolete
xfail marks.
Integration validation uses the companion [tirx-kernels
migration](https://github.com/mlc-ai/TIRx-kernels/pull/6), which updates
the registered kernels to the native builder API.
Validation on TVM `e17fcc98f9` plus this change:
- Incremental TVM build and pre-commit: passed.
- `python -m tirx_kernels.bench_suite --check-imports`: passed (274
workloads, 93 kernels).
- `python -m tirx_kernels.registry --strict`: passed.
- `python -m pytest tests/python/tirx/ -n 16`: 2926 passed, 308 skipped,
zero failures/errors/XPASS in 151.77 seconds on physical GPUs 0-4.
The skips retain the existing reference-dependency, toolchain, and
unsupported-configuration conditions.
---
python/tvm/backend/cuda/lang/clc.py | 4 +-
python/tvm/backend/cuda/lang/pipeline.py | 14 ++--
python/tvm/backend/cuda/lang/tile_scheduler.py | 2 +-
python/tvm/backend/cuda/tile_primitive/common.py | 14 ++--
.../backend/cuda/tile_primitive/copy/fallback.py | 19 ++++-
.../cuda/tile_primitive/copy/ld_stmatrix.py | 2 +-
.../cuda/tile_primitive/copy/vec_auto_gmem_smem.py | 2 +-
.../cuda/tile_primitive/copy/vec_auto_reg.py | 6 +-
.../backend/cuda/tile_primitive/copy/vec_forced.py | 4 +-
.../cuda/tile_primitive/copy_async/dsmem.py | 2 +
.../cuda/tile_primitive/elementwise/_common.py | 6 +-
.../backend/cuda/tile_primitive/elementwise/reg.py | 7 +-
.../cuda/tile_primitive/elementwise/smem.py | 2 +-
.../cuda/tile_primitive/gemm/mma_m16n8k_.py | 2 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 52 ++++---------
.../permute_layout/warp_xor_swizzle.py | 4 +-
.../backend/cuda/tile_primitive/reduction/local.py | 14 ++--
.../cuda/tile_primitive/reduction/shared.py | 14 ++--
src/tirx/op/builtin.cc | 13 ++--
src/tirx/transform/lower_tirx_cleanup.cc | 8 ++
tests/python/tirx/codegen/test_codegen_ampere.py | 12 +--
tests/python/tirx/codegen/test_codegen_cuda.py | 23 ++++--
tests/python/tirx/codegen/test_codegen_hopper.py | 14 +++-
tests/python/tirx/codegen/test_codegen_nvshmem.py | 2 +-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 6 +-
.../tile_primitive/cuda/copy/test_gmem_smem.py | 8 +-
.../tile_primitive/cuda/copy/test_ld_stmatrix.py | 8 +-
.../operator/tile_primitive/cuda/copy/test_reg.py | 67 +++++++++++++---
.../cuda/copy_async/test_tcgen05_cp.py | 6 +-
.../cuda/copy_async/test_tcgen05_ldst.py | 2 +-
.../tile_primitive/cuda/elementwise/test_binary.py | 80 +++++++++----------
.../tile_primitive/cuda/elementwise/test_unary.py | 72 ++++++++---------
.../cuda/gemm/test_gemm_mma_m16n8k_.py | 57 +++++++++-----
.../cuda/gemm_async/test_gemm_async.py | 2 +-
.../cuda/reduction/test_reduction.py | 90 +++++++++++-----------
tests/python/tirx/test_op_namespace_cleanup.py | 30 ++++++++
.../tirx/transform/test_transform_lower_tirx.py | 30 ++++++++
37 files changed, 426 insertions(+), 274 deletions(-)
diff --git a/python/tvm/backend/cuda/lang/clc.py
b/python/tvm/backend/cuda/lang/clc.py
index 670bc0c503..1b958cd240 100644
--- a/python/tvm/backend/cuda/lang/clc.py
+++ b/python/tvm/backend/cuda/lang/clc.py
@@ -39,7 +39,9 @@ def query_cancel_first_ctaid_x(first_ctaid_x, handle, *,
use_ld_acquire=True):
response = T.local_scalar("uint128")
canceled = T.local_scalar("uint32")
- T.ptx[f"ld{'.acquire.cta' if use_ld_acquire else
''}.shared.b128"](response, handle)
+ T.ptx[f"ld{'.acquire.cta' if T.constexpr(use_ld_acquire) else
''}.shared.b128"](
+ response, handle
+ )
T.ptx.clusterlaunchcontrol.query_cancel.is_canceled.pred.b128(canceled,
response)
T.tensor_store(first_ctaid_x.source, T.uint32(0xFFFFFFFF),
first_ctaid_x.indices)
T.ptx.clusterlaunchcontrol.query_cancel.get_first_ctaid__x.b32.b128(
diff --git a/python/tvm/backend/cuda/lang/pipeline.py
b/python/tvm/backend/cuda/lang/pipeline.py
index 5d21e30484..385fdf3007 100644
--- a/python/tvm/backend/cuda/lang/pipeline.py
+++ b/python/tvm/backend/cuda/lang/pipeline.py
@@ -73,7 +73,7 @@ class PipelineState:
@T.inline
def advance(self):
- if self.depth > 1:
+ if T.constexpr(self.depth > 1):
self.stage = self.stage + 1
if self.stage == self.depth:
self.stage = 0
@@ -219,11 +219,13 @@ class MBarrier:
@T.inline
def _arrive(self, bar, pred=None, count=None):
- if pred is None:
- T.ptx.mbarrier.arrive.shared.b64(bar, T.uint32(1 if count is None
else count))
+ if T.constexpr(pred is None):
+ T.ptx.mbarrier.arrive.shared.b64(
+ bar, T.uint32(1 if T.constexpr(count is None) else count)
+ )
else:
T.ptx.mbarrier.arrive.shared.b64(
- bar, T.uint32(1 if count is None else count), pred=pred
+ bar, T.uint32(1 if T.constexpr(count is None) else count),
pred=pred
)
def ptr_to(self, idx):
@@ -280,7 +282,7 @@ class TMABar(MBarrier):
@T.inline
def _arrive_tma_local(self, bar, tx_count=None):
- if tx_count is None:
+ if T.constexpr(tx_count is None):
T.ptx.mbarrier.arrive.shared.b64(bar, T.uint32(1))
else:
T.ptx.mbarrier.arrive.expect_tx.shared.b64(bar, T.uint32(tx_count))
@@ -307,7 +309,7 @@ class TCGen05Bar(MBarrier):
# form cannot depend on it, so a runtime mask always multicasts.
# ``pred`` rides the ptx keyword: the instruction is emitted
# predicated (@p) rather than branched around.
- if _tcgen05_commit_is_unicast(cta_mask):
+ if T.constexpr(_tcgen05_commit_is_unicast(cta_mask)):
T.evaluate(
T.ptx[
f"tcgen05.commit.cta_group::{cta_group}"
diff --git a/python/tvm/backend/cuda/lang/tile_scheduler.py
b/python/tvm/backend/cuda/lang/tile_scheduler.py
index f22902e20c..279b4524c2 100644
--- a/python/tvm/backend/cuda/lang/tile_scheduler.py
+++ b/python/tvm/backend/cuda/lang/tile_scheduler.py
@@ -821,7 +821,7 @@ class FlashAttentionLPTScheduler(BaseTileScheduler):
:class:`FlashAttentionLinearScheduler`, while keeping the LPT + L2
swizzle index mapping.
"""
- if self._num_ctas is None:
+ if T.constexpr(self._num_ctas is None):
self.linear_idx = self._total_tasks
else:
self.linear_idx = self.linear_idx + self._num_ctas
diff --git a/python/tvm/backend/cuda/tile_primitive/common.py
b/python/tvm/backend/cuda/tile_primitive/common.py
index 00372b8a26..c24f7e570b 100644
--- a/python/tvm/backend/cuda/tile_primitive/common.py
+++ b/python/tvm/backend/cuda/tile_primitive/common.py
@@ -209,18 +209,18 @@ def copy_vec_load_impl(
"""Implement copy operation with vectorized loads/stores."""
for s in T.serial(0, n_elements // (tx * vec_len)):
for tid_x in T.thread_binding(tx, "threadIdx.x"):
- if inst_type == CopyInstType.NORMAL:
+ if T.constexpr(inst_type == CopyInstType.NORMAL):
for vec in T.vectorized(vec_len):
fused = T.meta_var((s * tx + tid_x) * vec_len +
vec)
dst_indices = T.meta_var(get_indices(fused,
dst_st, dst_extent))
src_indices = T.meta_var(get_indices(fused,
src_st, src_extent))
dst[tuple(dst_indices)] = src[tuple(src_indices)]
- elif inst_type == CopyInstType.CP_ASYNC:
+ elif T.constexpr(inst_type == CopyInstType.CP_ASYNC):
fused = T.meta_var((s * tx + tid_x) * vec_len)
dst_indices = T.meta_var(get_indices(fused, dst_st,
dst_extent))
src_indices = T.meta_var(get_indices(fused, src_st,
src_extent))
- T.evaluate(T.ptx[f"cp.async.{'cg' if cp_size == 16
else 'ca'}.shared.global"](dst.ptr_to(dst_indices), src.ptr_to(src_indices),
cp_size)) # noqa: E501
- if dst.scope().startswith("shared") and inst_type ==
CopyInstType.NORMAL:
+ T.evaluate(T.ptx[f"cp.async.{'cg' if
T.constexpr(cp_size == 16) else 'ca'}.shared.global"](dst.ptr_to(dst_indices),
src.ptr_to(src_indices), cp_size)) # noqa: E501
+ if T.constexpr(dst.scope().startswith("shared") and inst_type ==
CopyInstType.NORMAL):
T.tvm_storage_sync("shared")
# fmt: on
elif sctx.is_thread:
@@ -228,17 +228,17 @@ def copy_vec_load_impl(
@T.function(check_well_formed=False)
def impl():
for s in T.serial(0, n_elements // (vec_len)):
- if inst_type == CopyInstType.NORMAL:
+ if T.constexpr(inst_type == CopyInstType.NORMAL):
for vec in T.vectorized(vec_len):
fused = T.meta_var(s * vec_len + vec)
dst_indices = T.meta_var(get_indices(fused, dst_st,
dst_extent))
src_indices = T.meta_var(get_indices(fused, src_st,
src_extent))
dst[tuple(dst_indices)] = src[tuple(src_indices)]
- elif inst_type == CopyInstType.CP_ASYNC:
+ elif T.constexpr(inst_type == CopyInstType.CP_ASYNC):
fused = T.meta_var(s * vec_len)
dst_indices = T.meta_var(get_indices(fused, dst_st,
dst_extent))
src_indices = T.meta_var(get_indices(fused, src_st,
src_extent))
- T.evaluate(T.ptx[f"cp.async.{'cg' if cp_size == 16 else
'ca'}.shared.global"](dst.ptr_to(dst_indices), src.ptr_to(src_indices),
cp_size)) # noqa: E501
+ T.evaluate(T.ptx[f"cp.async.{'cg' if T.constexpr(cp_size
== 16) else 'ca'}.shared.global"](dst.ptr_to(dst_indices),
src.ptr_to(src_indices), cp_size)) # noqa: E501
# fmt: on
else:
fail(f"unsupported exec_scope {sctx.scope_kind}")
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
index c2c0ce9793..d60169e48d 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
@@ -22,6 +22,7 @@ import warnings
import tvm
from tvm.script import tirx as T
from tvm.tirx import Function, Var
+from tvm.tirx.layout import ComposeLayout, TileLayout
from tvm.tirx.operator.tile_primitive.dispatcher import (
predicate,
register_dispatch,
@@ -39,6 +40,22 @@ def _region_st_extent(buffer_region):
return [r.min for r in region], [r.extent for r in region]
+def _is_scalar_copy(op_call: TilePrimitiveCall, sctx: DispatchContext):
+ ok, reason = _is_valid_copy(op_call, sctx)
+ if not ok:
+ return ok, reason
+ op_call = TilePrimitiveCall.downcast(op_call)
+ for region in (op_call.src, op_call.dst):
+ layout = region.source.layout
+ if isinstance(layout, ComposeLayout):
+ layout = layout.tile_layout
+ if isinstance(layout, TileLayout) and any(
+ it.axis.is_thread() for it in (*layout.shard, *layout.replica)
+ ):
+ return False, "single-thread fallback cannot access a tensor with
thread axes"
+ return True, None
+
+
def _emit_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext) ->
Function:
op_call = TilePrimitiveCall.downcast(op_call)
src: Var = op_call.src.source
@@ -114,7 +131,7 @@ def _emit_fallback(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Functio
"cuda",
variant="fallback",
priority=0,
- when=[predicate("validate_copy_op", _is_valid_copy)],
+ when=[predicate("scalar_copy_applicable", _is_scalar_copy)],
)
def copy_schedule_fallback(op_call: TilePrimitiveCall, sctx: DispatchContext)
-> Function:
return _emit_fallback(op_call, sctx)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
index c4c7eac1c5..73c6022ae9 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
@@ -376,7 +376,7 @@ def _emit(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Function:
r_words[r.apply(0, 0, 0, mm, i, 0,
shape=apply_shape)[r_mem_axis] // 2]
for i in range(num)
]
- if direction == "ld":
+ if T.constexpr(direction == "ld"):
T.ptx[ld_chain](*words, smem_ptr)
else:
# stmatrix reverses ldmatrix's operand order: address first.
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
index 246e7b5fb6..dee51b8315 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
@@ -176,7 +176,7 @@ def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Functi
s_off = s_apply_layout.apply(f, tid, v0, shape=apply_shape)["m"]
s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off)
g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin)
- if g_is_src:
+ if T.constexpr(g_is_src):
T.ptx[ld_g](*[tmp[i] for i in range(lanes)], g_ptr)
T.ptx[st_s](s_ptr, *[tmp[i] for i in range(lanes)])
else:
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
index 17412937b9..f57bc9d92b 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
@@ -599,7 +599,7 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Function:
@T.function(check_well_formed=False)
def impl():
- if not has_swizzle:
+ if T.constexpr(not has_swizzle):
s_base = _apply_s_layout(
s_apply_layout, thread_coords, 0, s_apply_shape, placeholders,
sctx
)
@@ -612,7 +612,7 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Function:
# replicates per-iter scratch arrays and pressures registers.
for f in range(total_outer):
ds, dr = _outer_const_offsets(outer, f)
- if has_swizzle:
+ if T.constexpr(has_swizzle):
s_off = _apply_s_layout(
s_apply_layout, thread_coords, f, s_apply_shape,
placeholders, sctx
)
@@ -623,7 +623,7 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Function:
# element pointer: the container index is the element offset scaled
# between element and word granularity.
r_w = (r_off_base + dr) * words_per_elem // elems_per_word
- if r_is_src:
+ if T.constexpr(r_is_src):
T.ptx[st_chain](s_ptr, *[r_words[r_w + i] for i in
range(lanes)])
else:
T.ptx[ld_chain](*[r_words[r_w + i] for i in range(lanes)],
s_ptr)
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 f3b72e9648..57afd2de2e 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
@@ -191,9 +191,9 @@ def _emit_forced_vec_copy(op_call: TilePrimitiveCall,
_sctx: DispatchContext, nu
# fmt: off
@T.function(check_well_formed=False)
def impl():
- if src_is_local:
+ if T.constexpr(src_is_local):
T.ptx[st_chain](dst_ptr, *_words(src, op_call.src))
- elif dst_is_local:
+ elif T.constexpr(dst_is_local):
T.ptx[ld_chain](*_words(dst, op_call.dst), src_ptr)
else:
tmp = T.alloc_local((lanes,), reg_dtype)
diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
index 041fbeb37e..a37f4845a1 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
@@ -134,6 +134,8 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx:
DispatchContext) -> Functi
# Helper to compute element offsets from loop variables (called via
T.meta_var)
def compute_offsets(loop_vars):
+ if isinstance(loop_vars, Var):
+ loop_vars = (loop_vars,)
src_off = 0
dst_off = 0
for j, v in enumerate(loop_vars):
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
index 90aea6abc5..21619d03fc 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
@@ -379,11 +379,11 @@ def emit_scope_sync(scope_kind: str):
@T.inline
def sync():
- if scope_kind == "cta":
+ if T.constexpr(scope_kind == "cta"):
T.cuda.cta_sync()
- elif scope_kind == "warpgroup":
+ elif T.constexpr(scope_kind == "warpgroup"):
T.cuda.warpgroup_sync(8)
- elif scope_kind == "warp":
+ elif T.constexpr(scope_kind == "warp"):
T.cuda.warp_sync()
return sync
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
index cc72a9d964..5854febc64 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
@@ -389,7 +389,7 @@ def _emit_induced_packed(
src_args = T.meta_var(
[
src.scalar
- if src.is_scalar
+ if T.constexpr(src.is_scalar)
else (
views[src.buf_region],
[[f * vec_len + k] for k in range(vec_len)],
@@ -423,7 +423,10 @@ def _emit_induced_scalar(
# Logical 1D coord = f (vec_len = 1 in scalar path); each
# buffer's layout maps to physical at access time.
src_vals = T.meta_var(
- [src.scalar if src.is_scalar else views[src.buf_region][f] for
src in srcs]
+ [
+ src.scalar if T.constexpr(src.is_scalar) else
views[src.buf_region][f]
+ for src in srcs
+ ]
)
views[dst_br][f] = T.cast(compute(src_vals, extras, dst_dtype),
dst_dtype)
diff --git a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
index 9b5178825c..82b2036feb 100644
--- a/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
@@ -212,7 +212,7 @@ def _emit_packed(plan, vec_impl, vec_chunk, total,
thread_cnt, sctx) -> Function
src_args = T.meta_var(
[
srcs[i].scalar
- if srcs[i].is_scalar
+ if T.constexpr(srcs[i].is_scalar)
else (
srcs[i].buf_region.source,
_src_lane_indices(
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
index 8e28a2e712..ed2c50d0cd 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
@@ -570,7 +570,7 @@ def gemm_cuda_mma_dispatch(op: TilePrimitiveCall, sctx:
DispatchContext) -> Func
# Initialize D[m, n]: copy C (beta==1) or clear to 0 (beta==0).
for rM in T.unroll(n_rM):
for rN in T.unroll(n_rN):
- if use_c:
+ if T.constexpr(use_c):
d_local[m, n, rM, rN] = c_local[m, n, rM, rN]
else:
d_local[m, n, rM, rN] = T.float32(0)
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 f3e43707bc..de5b7983b1 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -1290,6 +1290,13 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
# and uses rows 64-127 for the other half.
N_mma_phys_cols = N_mma // 2 if is_2x2 or packed_n2 else N_mma
+ def _d_operand(mi, ni):
+ # Simplify coordinates before TVMScript binds them to IR variables, so
+ # constant lane offsets remain visible to the address fast path.
+ row = tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma)
+ col = tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b)
+ return _get_tmem_addr_fast(tmem_addr, row, col)
+
# The ptx instruction spelling, resolved once from the trace-time dtypes.
if is_block_scaled:
_bs_kind = _get_tcgen05_mma_kind(C_type, A_type, B_type, SFA_type,
SFB_type)
@@ -1326,7 +1333,6 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
for mi in T.unroll(M_tiles):
for ni in T.unroll(N_tiles):
for ki in T.unroll(K_iters):
- # meta_var inlines operands into mma.block_scale (avoids
LMEM temps).
a_val = T.meta_var(_a_operand(mi, ki, descA_in))
descB_val = T.meta_var(_b_desc_val(descB_in, ni, ki))
should_accum = T.meta_var(tvm.tirx.any(ki != 0,
accum_expr))
@@ -1353,32 +1359,18 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
tvm.tirx.floordiv(sfb_tcol, SFB_elem_per_col),
)
)
- if needs_sf_id:
+ if T.constexpr(needs_sf_id):
sf_id =
T.meta_var(analyzer.simplify(tvm.tirx.floormod(sfa_tcol, SFA_elem_per_col))) #
noqa: E501
T.cuda.runtime_instr_desc(T.address_of(descI_in),
sf_id)
- tmem_col = T.meta_var(
- tmem_offset_32b + ni * (N_mma_phys_cols //
C_elem_per_32b)
- )
- tmem_row = T.meta_var(
- tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma)
- )
if elect_pred:
T.evaluate(T.ptx[mma_chain](
- T.cast(_get_tmem_addr_fast(tmem_addr, tmem_row,
tmem_col), "uint32"),
- T.cast(a_val, "uint32") if a_is_tmem else a_val,
+ T.cast(_d_operand(mi, ni), "uint32"),
+ T.cast(a_val, "uint32") if T.constexpr(a_is_tmem)
else a_val,
descB_val, descI_in,
T.cast(sfa_addr, "uint32"), T.cast(sfb_addr,
"uint32"),
should_accum,
))
else:
- # Wrap each per-MMA operand in ``T.meta_var`` so the parser inlines
- # the value directly into the tcgen05.mma call instead of
- # materializing it into a fresh ``alignas(64) T x[1]; x[0] = expr``
- # local. Without this wrap each unrolled MMA emits 4 throw-away
- # 1-element local arrays (``a_val_ptr``, ``descB_val_ptr``,
- # ``should_accum_ptr``, ``tmem_col_ptr``) which ptxas cannot fold
- # back into the operand and the resulting LMEM round-trips show up
- # on the fa4 hot path.
if encode_per_mma:
@T.inline
def main_impl(descA_in, descB_in, descI_in):
@@ -1386,22 +1378,16 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
for ni in T.unroll(N_tiles):
for ki in T.unroll(K_iters):
should_accum = T.meta_var(tvm.tirx.any(ki != 0,
accum_expr))
- tmem_col = T.meta_var(
- tmem_offset_32b + ni * (N_mma_phys_cols //
C_elem_per_32b)
- )
- tmem_row = T.meta_var(
- tmem_lane_offset + (0 if M_tiles == 1 else mi *
M_mma)
- )
if elect_pred:
descB_mma = T.meta_var(
_encoded_desc_val(
B_buffer, _b_offset(ni, ki), B_ldo, B_sdo,
B_swizzle_mode.value
)
)
- if a_is_tmem:
+ if T.constexpr(a_is_tmem):
a_val = T.meta_var(_a_operand(mi, ki,
descA_in))
_emit_mma(
- _get_tmem_addr_fast(tmem_addr, tmem_row,
tmem_col),
+ _d_operand(mi, ni),
a_val, descB_mma, descI_in, should_accum,
)
else:
@@ -1415,7 +1401,7 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
)
)
_emit_mma(
- _get_tmem_addr_fast(tmem_addr, tmem_row,
tmem_col),
+ _d_operand(mi, ni),
descA_mma, descB_mma, descI_in,
should_accum,
)
else:
@@ -1427,15 +1413,9 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
a_val = T.meta_var(_a_operand(mi, ki, descA_in))
descB_val = T.meta_var(_b_desc_val(descB_in, ni, ki))
should_accum = T.meta_var(tvm.tirx.any(ki != 0,
accum_expr))
- tmem_col = T.meta_var(
- tmem_offset_32b + ni * (N_mma_phys_cols //
C_elem_per_32b)
- )
- tmem_row = T.meta_var(
- tmem_lane_offset + (0 if M_tiles == 1 else mi *
M_mma)
- )
if elect_pred:
_emit_mma(
- _get_tmem_addr_fast(tmem_addr, tmem_row,
tmem_col),
+ _d_operand(mi, ni),
a_val, descB_val, descI_in, should_accum,
)
@@ -1443,7 +1423,7 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
@T.inline
def call_main(descI_arg):
- if local_hoist:
+ if T.constexpr(local_hoist):
descB_local = T.alloc_local((1,), "uint64")
T.cuda.tcgen05.encode_matrix_descriptor(
T.address_of(descB_local[0]),
@@ -1454,7 +1434,7 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
)
# local_hoist runs at the call site under elected-thread control,
so the
# descriptor is consumed by the same thread and needs no warp
shuffle.
- if A_use_add:
+ if T.constexpr(A_use_add):
descA_local = T.alloc_local((1,), "uint64")
T.cuda.tcgen05.encode_matrix_descriptor(
T.address_of(descA_local[0]),
diff --git
a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
index 36787460da..e141c1e9fe 100644
--- a/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
+++ b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
@@ -328,7 +328,7 @@ def _impl(op_call, sctx):
# fmt: off
if direct:
- @T.function
+ @T.function(check_well_formed=False)
def impl():
warp_size = T.meta_var(32)
lane_id = T.meta_var(tid_x % warp_size)
@@ -354,7 +354,7 @@ def _impl(op_call, sctx):
T.ptx[st_chain](ptr, regs[r])
T.cuda.warp_sync()
else:
- @T.function
+ @T.function(check_well_formed=False)
def impl():
warp_size = T.meta_var(32)
lane_id = T.meta_var(tid_x % warp_size)
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/local.py
b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
index a328757aac..feda434141 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/local.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
@@ -141,7 +141,7 @@ def _gen_warp_shuffle_reduce(src, dst, reduce_width,
local_elems, accum, op_type
src_local = src.local(local_elems, layout=src.layout.storage())
dst_local = dst.local(local_elems, layout=dst.layout.storage())
for k in T.serial(local_elems):
- if not is_same_buffer:
+ if T.constexpr(not is_same_buffer):
dst_local[k] = src_local[k]
dst_local[k] = T.cuda.warp_reduce(dst_local[k], op_str,
reduce_width)
# fmt: on
@@ -273,7 +273,7 @@ def _emit_reduction_local_thread_wise(
def impl():
for spa in T.serial(spatial_len):
dst_idx = T.meta_var(get_indices(spa, dst_st, dst_extent))
- if not accum:
+ if T.constexpr(not accum):
dst[tuple(dst_idx)] = init_value
for red in T.serial(reduction_len):
src_idx = T.meta_var(get_src_indices(spa, red))
@@ -367,12 +367,12 @@ def _emit_reduction_local_view(
for spa in T.serial(dst_local_total):
dst_idx = T.meta_var(get_indices(spa, dst_local_st,
dst_local_ext))
old_val[0] = dst_local[tuple(dst_idx)]
- if not in_place:
+ if T.constexpr(not in_place):
dst_local[tuple(dst_idx)] = init_value
for red in T.serial(reduction_local_total):
src_idx = T.meta_var(_get_src_local_index(spa, red))
dst_local[tuple(dst_idx)] =
op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
- if shuffle:
+ if T.constexpr(shuffle):
mask = T.tvm_warp_activemask()
shuffle_data(mask, dst_local, dst_idx)
dst_local[tuple(dst_idx)] = op_func(dst_local[tuple(dst_idx)],
old_val[0])
@@ -386,13 +386,13 @@ def _emit_reduction_local_view(
for spa in T.serial(dst_local_total):
dst_idx = T.meta_var(get_indices(spa, dst_local_st,
dst_local_ext))
- if not in_place:
- if not accum:
+ if T.constexpr(not in_place):
+ if T.constexpr(not accum):
dst_local[tuple(dst_idx)] = init_value
for red in T.serial(reduction_local_total):
src_idx = T.meta_var(_get_src_local_index(spa, red))
dst_local[tuple(dst_idx)] =
op_func(dst_local[tuple(dst_idx)], src_local[tuple(src_idx)]) # noqa: E501
- if shuffle:
+ if T.constexpr(shuffle):
mask = T.tvm_warp_activemask()
shuffle_data(mask, dst_local, dst_idx)
# fmt: on
diff --git a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
index 72c3a46da3..cfcf135a4d 100644
--- a/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
+++ b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
@@ -179,17 +179,17 @@ def _emit_reduction_shared_cta(
@T.inline
def sync():
- if exec_scope_name == "cta":
+ if T.constexpr(exec_scope_name == "cta"):
T.cuda.cta_sync()
- elif exec_scope_name == "warpgroup":
+ elif T.constexpr(exec_scope_name == "warpgroup"):
T.cuda.warpgroup_sync(8) # TODO: fix this hardcoded value
- elif exec_scope_name == "warp":
+ elif T.constexpr(exec_scope_name == "warp"):
T.cuda.warp_sync()
- elif exec_scope_name == "thread":
+ elif T.constexpr(exec_scope_name == "thread"):
pass
# fmt: off
- @T.function
+ @T.function(check_well_formed=False)
def impl():
tid_in_scope = get_tid_in_scope()
thread_data = T.alloc_tensor([1], dtype=dtype, scope="local")
@@ -237,11 +237,11 @@ def _emit_reduction_shared_thread(
assert op_func is not None
init_value = reduce_default_value_table(dtype).get(reduce_op)
- @T.function
+ @T.function(check_well_formed=False)
def impl():
for spa_fused in T.serial(spatial_len):
dst_indices = T.meta_var(get_indices(spa_fused, dst_st,
dst_extent))
- if not accum:
+ if T.constexpr(not accum):
dst[tuple(dst_indices)] = init_value
for red_fused in T.serial(reduction_len):
src_indices = T.meta_var(
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index ee755d37ef..38cdd0d866 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -37,14 +37,15 @@ void CallFFIKernelAttr::RegisterReflection() {
}
void TensorMapEncodeTiledAttr::RegisterReflection() {
+ namespace refl = ffi::reflection;
ffi::reflection::ObjectDef<TensorMapEncodeTiledAttr>()
.def_ro("descriptor_dtype", &TensorMapEncodeTiledAttr::descriptor_dtype)
.def_ro("rank", &TensorMapEncodeTiledAttr::rank)
- .def_ro("interleave", &TensorMapEncodeTiledAttr::interleave)
- .def_ro("swizzle", &TensorMapEncodeTiledAttr::swizzle)
- .def_ro("l2_promotion", &TensorMapEncodeTiledAttr::l2_promotion)
- .def_ro("oob_fill", &TensorMapEncodeTiledAttr::oob_fill)
- .def_ro("force_cu_dtype", &TensorMapEncodeTiledAttr::force_cu_dtype);
+ .def_ro("interleave", &TensorMapEncodeTiledAttr::interleave,
refl::DefaultValue(0))
+ .def_ro("swizzle", &TensorMapEncodeTiledAttr::swizzle,
refl::DefaultValue(0))
+ .def_ro("l2_promotion", &TensorMapEncodeTiledAttr::l2_promotion,
refl::DefaultValue(0))
+ .def_ro("oob_fill", &TensorMapEncodeTiledAttr::oob_fill,
refl::DefaultValue(0))
+ .def_ro("force_cu_dtype", &TensorMapEncodeTiledAttr::force_cu_dtype,
refl::DefaultValue(-1));
}
ffi::Array<Var> LaunchThreadBodyParams(const CallNode* call) {
@@ -567,7 +568,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
OpDef("tirx.tensormap_encode_tiled")
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(32))
.signature(sig::arg("descriptor", "The descriptor."), sig::arg("data",
"The input data."),
- sig::var_args("args"))
+ sig::var_args("args"),
sig::call_attrs<TensorMapEncodeTiledAttr>())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.tensormap_encode_tiled"))
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc
b/src/tirx/transform/lower_tirx_cleanup.cc
index afab85154c..b71d58b719 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -235,6 +235,14 @@ class LayoutApplier : public IRMutatorWithAnalyzer {
.ValueOrUnchanged(type->strides[i])
.as_or_throw<PrimExpr>());
}
+ // TMEM addresses may load from an allocation whose variable was rebuilt
+ // above. Keep the type metadata in sync with the declaration's pointer.
+ for (size_t i = 0; i < type->allocated_addr.size(); ++i) {
+ type->allocated_addr.Set(
+ i, StmtExprMutator::Mutate(ffi::AnyView(type->allocated_addr[i]),
InplaceMode::kDisallow)
+ .ValueOrUnchanged(type->allocated_addr[i])
+ .as_or_throw<PrimExpr>());
+ }
type->layout = std::nullopt;
type->elem_offset =
StmtExprMutator::Mutate(ffi::AnyView(buf->elem_offset),
InplaceMode::kDisallow)
diff --git a/tests/python/tirx/codegen/test_codegen_ampere.py
b/tests/python/tirx/codegen/test_codegen_ampere.py
index 754a026aa2..0fea3668a9 100644
--- a/tests/python/tirx/codegen/test_codegen_ampere.py
+++ b/tests/python/tirx/codegen/test_codegen_ampere.py
@@ -105,13 +105,13 @@ def test_ptx_mma_m16n8k16(a_type, no_c_ptr):
@T.inline
def G2L(buf_local, buf_global, block_8x8, mode="row"):
- if mode == "row":
+ if T.constexpr(mode == "row"):
for i in range(block_8x8):
row = T.meta_var(i % 2 * 8 + tx // 4)
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
for j in range(2):
buf_local[i * 2 + j] = buf_global[row, col + j]
- elif mode == "col":
+ elif T.constexpr(mode == "col"):
for i in range(block_8x8):
row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
col = T.meta_var(i // 2 * 8 + tx // 4)
@@ -129,7 +129,7 @@ def test_ptx_mma_m16n8k16(a_type, no_c_ptr):
B_words = B_local.view("uint32")
# ptx takes the accumulator as an operand; the legacy "omit c" form
# fed literal zeros, which is now spelled at the call site.
- if no_c_ptr:
+ if T.constexpr(no_c_ptr):
for i in range(4):
C_local[i] = T.float32(0)
T.ptx[f"mma.sync.aligned.m16n8k16.row.col.f32.{_elem}.{_elem}.f32"](
@@ -182,13 +182,13 @@ def test_ptx_mma_m16n8k8(a_type, no_c_ptr):
@T.inline
def G2L(buf_local, buf_global, block_8x8, mode="row"):
- if mode == "row":
+ if T.constexpr(mode == "row"):
for i in range(block_8x8):
row = T.meta_var(i % 2 * 8 + tx // 4)
col = T.meta_var(i // 2 * 8 + (tx % 4) * 2)
for j in range(2):
buf_local[i * 2 + j] = buf_global[row, col + j]
- elif mode == "col":
+ elif T.constexpr(mode == "col"):
for i in range(block_8x8):
row = T.meta_var(i % 2 * 8 + (tx % 4) * 2)
col = T.meta_var(i // 2 * 8 + tx // 4)
@@ -205,7 +205,7 @@ def test_ptx_mma_m16n8k8(a_type, no_c_ptr):
B_words = B_local.view("uint32")
# ptx takes the accumulator as an operand; the legacy "omit c" form
# fed literal zeros, which is now spelled at the call site.
- if no_c_ptr:
+ if T.constexpr(no_c_ptr):
for i in range(4):
C_local[i] = T.float32(0)
T.ptx[f"mma.sync.aligned.m16n8k8.row.col.f32.{_elem}.{_elem}.f32"](
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index 21aa7428a6..9a782adc1c 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -1108,9 +1108,9 @@ def test_ptx_cp_async(cp_size, cache_hint, prefetch_size,
predicate, fill_mode):
for i in T.vectorized(N):
A_shared[i] = 5.0
T.ptx.fence.proxy.async_.shared__cta()
- if fill_mode == "zero":
+ if T.constexpr(fill_mode == "zero"):
T.ptx[chain](A_shared.ptr_to([0]), A.ptr_to([0]), cp_size,
src_size, *cache_args)
- elif has_pred:
+ elif T.constexpr(has_pred):
T.ptx[chain](A_shared.ptr_to([0]), A.ptr_to([0]), cp_size,
*cache_args, pred=predicate)
else:
T.ptx[chain](A_shared.ptr_to([0]), A.ptr_to([0]), cp_size,
*cache_args)
@@ -1161,19 +1161,28 @@ def test_ptx_ldmatrix(trans, num):
# contiguous fp16[8] buffer, so the registers land through a uint32
# view, two fp16 elements per word.
A_words = A_local.view("uint32")
- if num == 1:
- T.ptx[f"ldmatrix.sync.aligned.m8n8.x1{'.trans' if trans else
''}.shared.b16"](
+ if T.constexpr(num == 1):
+ T.ptx[
+ "ldmatrix.sync.aligned.m8n8.x1"
+ f"{'.trans' if T.constexpr(trans) else ''}.shared.b16"
+ ](
A_words[0],
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
)
- elif num == 2:
- T.ptx[f"ldmatrix.sync.aligned.m8n8.x2{'.trans' if trans else
''}.shared.b16"](
+ elif T.constexpr(num == 2):
+ T.ptx[
+ "ldmatrix.sync.aligned.m8n8.x2"
+ f"{'.trans' if T.constexpr(trans) else ''}.shared.b16"
+ ](
A_words[0],
A_words[1],
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
)
else:
- T.ptx[f"ldmatrix.sync.aligned.m8n8.x4{'.trans' if trans else
''}.shared.b16"](
+ T.ptx[
+ "ldmatrix.sync.aligned.m8n8.x4"
+ f"{'.trans' if T.constexpr(trans) else ''}.shared.b16"
+ ](
A_words[0],
A_words[1],
A_words[2],
diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py
b/tests/python/tirx/codegen/test_codegen_hopper.py
index f8e4dc467f..ad9d48bec6 100644
--- a/tests/python/tirx/codegen/test_codegen_hopper.py
+++ b/tests/python/tirx/codegen/test_codegen_hopper.py
@@ -69,7 +69,7 @@ def test_ptx_setmaxnreg(inc):
T.device_entry()
cta_id = T.cta_id([1])
tid = T.thread_id([128])
- T.ptx[f"setmaxnreg.{'inc' if inc else 'dec'}.sync.aligned.u32"](32)
+ T.ptx[f"setmaxnreg.{'inc' if T.constexpr(inc) else
'dec'}.sync.aligned.u32"](32)
# fmt: on
src, mod = _get_source(func)
@@ -97,7 +97,7 @@ def test_stmatrix_sync_aligned(trans):
# stmatrix stores 4 b32 registers; reg is fp16, so they ride a uint32
# view, two elements per word.
reg_words = reg.view("uint32")
- T.ptx[f"stmatrix.sync.aligned.m8n8.x4{'.trans' if trans else
''}.shared.b16"](
+ T.ptx[f"stmatrix.sync.aligned.m8n8.x4{'.trans' if T.constexpr(trans)
else ''}.shared.b16"](
A_smem.ptr_to([tx % 16, tx // 16 * 8]),
reg_words[0], reg_words[1], reg_words[2], reg_words[3],
)
@@ -166,7 +166,10 @@ def test_ptx_stmatrix(trans, num):
for i in range(8):
A_local[i] = (i // 2) * 64 + tx * 2 + i % 2
A_words = A_local.view("uint32")
- T.ptx[f"stmatrix.sync.aligned.m8n8.x{num}{'.trans' if trans else
''}.shared.b16"](
+ T.ptx[
+ f"stmatrix.sync.aligned.m8n8.x{num}"
+ f"{'.trans' if T.constexpr(trans) else ''}.shared.b16"
+ ](
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
*[A_words[i] for i in range(num)],
)
@@ -239,7 +242,10 @@ def test_ptx_stmatrix_noncontiguous(trans, num):
A_local[i * STRIDE + 0] = T.float16(i * 64 + tx * 2 + 0)
A_local[i * STRIDE + 1] = T.float16(i * 64 + tx * 2 + 1)
A_words = A_local.view("uint32")
- T.ptx[f"stmatrix.sync.aligned.m8n8.x{num}{'.trans' if trans else
''}.shared.b16"](
+ T.ptx[
+ f"stmatrix.sync.aligned.m8n8.x{num}"
+ f"{'.trans' if T.constexpr(trans) else ''}.shared.b16"
+ ](
A_shared.ptr_to([tx % 16, tx // 16 * 8]),
*[A_words[i * STRIDE // 2] for i in range(num)],
)
diff --git a/tests/python/tirx/codegen/test_codegen_nvshmem.py
b/tests/python/tirx/codegen/test_codegen_nvshmem.py
index 16d82cb9fc..1516ef199e 100644
--- a/tests/python/tirx/codegen/test_codegen_nvshmem.py
+++ b/tests/python/tirx/codegen/test_codegen_nvshmem.py
@@ -143,7 +143,7 @@ def test_codegen_nvshmem():
my_pe = T.nvshmem.my_pe()
n_pes = T.nvshmem.n_pes()
dst_pe = (my_pe + 1) % n_pes
- if sig_op == "add":
+ if T.constexpr(sig_op == "add"):
res[0] = 1
T.nvshmem.barrier_all()
T.nvshmem.signal_op(sig_addr=res.ptr_to([0]), signal=1,
sig_op=sig_op, pe=dst_pe)
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
index 5f12ba5059..036c171380 100644
--- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
+++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
@@ -83,13 +83,13 @@ def _shared_scratch_copy_kernel(num_bytes: int):
src_buf = T.alloc_tensor((nelems,), smem_dtype, scope="shared")
dst_buf = T.alloc_tensor((nelems,), smem_dtype, scope="shared")
tmp = T.alloc_local((lanes,), reg_dtype)
- if fill_offset is not None:
+ if T.constexpr(fill_offset is not None):
if lane < nelems:
src_buf[lane] = T.uint32(lane + fill_offset)
- elif fill_fp16 is not None:
+ elif T.constexpr(fill_fp16 is not None):
if lane == 0:
src_buf[0] = T.float16(fill_fp16)
- elif fill_u8 is not None:
+ elif T.constexpr(fill_u8 is not None):
if lane == 0:
src_buf[0] = T.uint8(fill_u8)
elif lane == 0:
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
index 63cad68de1..7c73261729 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
@@ -275,10 +275,6 @@ def _align(
)
[email protected](
- reason="align_layouts_gs ignores swizzle chunk size; "
- "_extract_tile strips the swizzle wrap before vec_len pick."
-)
@pytest.mark.parametrize("per_element,expected_max_vec", [(2, 4), (1, 2), (0,
1)])
def test_swizzled_smem_vec_len_must_fit_chunk(per_element, expected_max_vec):
"""A swizzled ``ComposeLayout`` keeps the bottom ``per_element``
@@ -538,9 +534,7 @@ def test_gmem_smem_swizzle_uses_structured_compose_apply():
src = ex.mod.imports[0].inspect_source()
s_off_lines = [
- line
- for line in src.splitlines()
- if line.strip().startswith("s_off_ptr") and "[0] =" in line
+ line for line in src.splitlines() if line.strip().startswith("int
s_off") and " = " in line
]
assert len(s_off_lines) == 2, "expected one structured S offset in each
copy direction"
assert all("^" in line for line in s_off_lines)
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
index 6a359dc275..2dee76532d 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
@@ -428,7 +428,9 @@ def test_ldstmatrix_swizzle_multi_iter_pow2():
assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
address_lines = [
- line for line in src.splitlines() if "smem_off_ptr" in line and "[0]
=" in line
+ line
+ for line in src.splitlines()
+ if line.strip().startswith("int smem_off") and " = " in line
]
assert len(address_lines) == 8
assert all("& 7) << 3" in line for line in address_lines)
@@ -499,7 +501,9 @@ def test_ldstmatrix_swizzle_multi_iter_linear():
assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src
address_lines = [
- line for line in src.splitlines() if "smem_off_ptr" in line and "[0]
=" in line
+ line
+ for line in src.splitlines()
+ if line.strip().startswith("int smem_off") and " = " in line
]
assert len(address_lines) == 10
assert all("& 7) << 3" in line for line in address_lines)
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
index 483110858c..93eb2caf93 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
@@ -458,17 +458,43 @@ def test_reg_copy_linear_shared_hoists_thread_base():
T.cuda.cta_sync()
T.evaluate(smem[tid, 0])
+ dispatched = []
+
+ @tvm.instrument.pass_instrument
+ class CaptureDispatch:
+ def run_after_pass(self, mod, info):
+ if info.name == "tirx.TilePrimitiveDispatch":
+ dispatched.append(mod["main"])
+
target = tvm.target.Target("cuda")
- with target:
+ with target, tvm.transform.PassContext(instruments=[CaptureDispatch()]):
ex = tvm.compile(tvm.IRModule({"main": kernel}), target=target,
tir_pipeline="tirx")
src = ex.mod.imports[0].inspect_source()
- base_assignment = "s_base_ptr[0] ="
- loop = "for (int f = 0; f < 8; ++f)"
- assert base_assignment in src and loop in src
- assert src.index(base_assignment) < src.index(loop)
- assert "(s_base_ptr[0] + ds_ptr[0])" in src
- assert "s_off_ptr" not in src
+ nodes = []
+ tvm_ffi.structural_walk(dispatched[0].body, nodes.append)
+ (copy_loop,) = [
+ node for node in nodes if isinstance(node, tvm.tirx.For) and
int(node.extent) == 8
+ ]
+ loop_nodes = []
+ tvm_ffi.structural_walk(copy_loop.body, loop_nodes.append)
+ (pointer,) = [
+ node
+ for node in loop_nodes
+ if isinstance(node, tvm.ir.Call)
+ and node.op.name == "tirx.cuda.func_call"
+ and node.args[0].value == "tvm_builtin_pointer_offset"
+ ]
+ offset = pointer.args[2]
+ assert isinstance(offset, tvm.tirx.Add)
+ (base,) = [
+ node for node in nodes if isinstance(node, tvm.tirx.Bind) and
node.var.same_as(offset.a)
+ ]
+ assert not any(isinstance(node, tvm.tirx.Bind) and node.same_as(base) for
node in loop_nodes)
+ assert not any(
+ var.same_as(copy_loop.loop_var) for var in
tvm.tirx.analysis.undefined_vars(base.value)
+ )
+ assert "st.shared.v4.u32" in src
@pytest.mark.gpu
@@ -536,9 +562,7 @@ def
test_reg_copy_wg_local_to_swizzled_shared_uses_structured_compose_apply():
# (2) Structured address fingerprint: tid contributes one atom-aligned
# add, while the bounded outer coordinate is XORed with the phase.
s_off_lines = [
- line
- for line in src.splitlines()
- if line.strip().startswith("s_off_ptr") and "[0] =" in line
+ line for line in src.splitlines() if line.strip().startswith("int
s_off") and " = " in line
]
assert len(s_off_lines) == 1
assert "^" in s_off_lines[0]
@@ -605,6 +629,29 @@ def test_copy_fallback_handles_scalar_regions():
assert "dst_ptr[2] = src_ptr[0];" in src
[email protected]("swizzled", [False, True])
+def test_copy_fallback_rejects_distributed_tensors(swizzled):
+ from tvm.tirx.layout import ComposeLayout
+
+ shape = (32, 4)
+ layout = TileLayout(S[shape : (1 @ tx, 1)])
+ if swizzled:
+ layout = ComposeLayout(0, 1, 2, layout)
+
+ @T.function
+ def kernel():
+ T.device_entry()
+ T.cta_id([1])
+ T.thread_id([32])
+ reg = T.alloc_tensor(shape, "float32", scope="local", layout=layout)
+ smem = T.alloc_tensor(shape, "float32", scope="shared")
+ Tx.cta.copy(smem, reg, dispatch="fallback")
+
+ with tvm.target.Target("cuda"):
+ with pytest.raises(RuntimeError, match="single-thread fallback cannot
access"):
+ tvm.tirx.transform.LowerTIRx()(tvm.IRModule({"main": kernel}))
+
+
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >=
9.0")
@pytest.mark.parametrize(
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
index 1b7f72c4d9..fc89f2672b 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
@@ -173,7 +173,7 @@ def _make_cp_kernel(
tmem = T.decl_tensor(
t_full_shape, dtype, scope="tmem",
allocated_addr=tmem_addr[0], layout=t_full
)
- if pre_zero:
+ if T.constexpr(pre_zero):
zero_reg = T.alloc_tensor((W32,), "uint32", scope="local")
for i in range(W32):
zero_reg[i] = T.uint32(0)
@@ -384,7 +384,7 @@ def test_cp_shape_config_routes_to_generic_planner():
src = mod.mod.imports[0].inspect_source()
assert "tcgen05.cp.cta_group::1.128x256b" in src
# Generic-planner signature: descriptor template encoded at address 0...
- assert "reinterpret_cast<void*>((uint64_t)0)" in src
+ assert "reinterpret_cast<void*>((uint64_t)0ULL)" in src
# ...and patched per cp via the 0x3FFF address-field mask.
assert src.count("cp_desc_ptr[0] &") == 4, f"expected 4 patched cps;
src=\n{src}"
@@ -678,7 +678,7 @@ def
test_cp_default_32x128b_instruction_sequence_unchanged():
]
assert len(encode_lines) == 1, encode_lines
# (ldo, sdo, swizzle) = (0, 8, 0), template encoded at smem address 0.
- assert "reinterpret_cast<void*>((uint64_t)0), 0, 8, 0);" in
encode_lines[0], encode_lines[0]
+ assert "reinterpret_cast<void*>((uint64_t)0ULL), 0, 8, 0);" in
encode_lines[0], encode_lines[0]
cp_lines = [
line
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py
index b355722146..6e7f3f00e2 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py
@@ -742,7 +742,7 @@ def test_datapath_B_codegen(direction):
layout=tmem_datapath_layout("B", 64, n_cols),
)
frag = T.alloc_tcgen05_ldst_frag("32x32b", (64, n_cols), "float32")
- if direction == "ld":
+ if T.constexpr(direction == "ld"):
Tx.wg.copy_async(frag[:, :], tmem[:, :])
else:
Tx.wg.copy_async(tmem[:, :], frag[:, :])
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
index 13c1abb586..4f5894af17 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py
@@ -103,19 +103,19 @@ def test_binary_op_shared(input, op_type, operands_type,
dtype):
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
Tx.cta.copy(B_smem[tuple(copy_slice)], B[tuple(copy_slice)])
T.cuda.cta_sync()
- if op_type == "add":
+ if T.constexpr(op_type == "add"):
Tx.cta.add(
A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)],
B_smem[tuple(map_slice_b)]
)
- elif op_type == "sub":
+ elif T.constexpr(op_type == "sub"):
Tx.cta.sub(
A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)],
B_smem[tuple(map_slice_b)]
)
- elif op_type == "mul":
+ elif T.constexpr(op_type == "mul"):
Tx.cta.mul(
A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)],
B_smem[tuple(map_slice_b)]
)
- elif op_type == "fdiv":
+ elif T.constexpr(op_type == "fdiv"):
Tx.cta.fdiv(
A_smem[tuple(map_slice_res)], A_smem[tuple(map_slice_a)],
B_smem[tuple(map_slice_b)]
)
@@ -134,25 +134,25 @@ def test_binary_op_shared(input, op_type, operands_type,
dtype):
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
T.cuda.cta_sync()
- if op_type == "add":
- if operands_type == "const_region":
+ if T.constexpr(op_type == "add"):
+ if T.constexpr(operands_type == "const_region"):
Tx.cta.add(A_smem[tuple(map_slice_res)], const,
A_smem[tuple(map_slice_a)])
- elif operands_type == "region_const":
+ elif T.constexpr(operands_type == "region_const"):
Tx.cta.add(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)], const)
- elif op_type == "sub":
- if operands_type == "const_region":
+ elif T.constexpr(op_type == "sub"):
+ if T.constexpr(operands_type == "const_region"):
Tx.cta.sub(A_smem[tuple(map_slice_res)], const,
A_smem[tuple(map_slice_a)])
- elif operands_type == "region_const":
+ elif T.constexpr(operands_type == "region_const"):
Tx.cta.sub(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)], const)
- elif op_type == "mul":
- if operands_type == "const_region":
+ elif T.constexpr(op_type == "mul"):
+ if T.constexpr(operands_type == "const_region"):
Tx.cta.mul(A_smem[tuple(map_slice_res)], const,
A_smem[tuple(map_slice_a)])
- elif operands_type == "region_const":
+ elif T.constexpr(operands_type == "region_const"):
Tx.cta.mul(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)], const)
- elif op_type == "fdiv":
- if operands_type == "const_region":
+ elif T.constexpr(op_type == "fdiv"):
+ if T.constexpr(operands_type == "const_region"):
Tx.cta.fdiv(A_smem[tuple(map_slice_res)], const,
A_smem[tuple(map_slice_a)])
- elif operands_type == "region_const":
+ elif T.constexpr(operands_type == "region_const"):
Tx.cta.fdiv(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)], const)
T.cuda.cta_sync()
Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
@@ -226,9 +226,9 @@ def test_binary_non_commutative_const_lhs_rejected(op_type):
_bx = T.cta_id([1])
_tid = T.thread_id([64])
A_smem = T.alloc_tensor(shape, dtype, scope="shared",
layout=layout)
- if op_type == "sub":
+ if T.constexpr(op_type == "sub"):
Tx.cta.sub(A_smem, const, A_smem)
- elif op_type == "fdiv":
+ elif T.constexpr(op_type == "fdiv"):
Tx.cta.fdiv(A_smem, const, A_smem)
target = tvm.target.Target("cuda")
@@ -268,10 +268,10 @@ def test_binary_op_shared_subcta_scope(exec_scope,
op_type):
Tx.cta.copy(A_smem, A)
Tx.cta.copy(B_smem, B)
T.cuda.cta_sync()
- if exec_scope == "warp":
+ if T.constexpr(exec_scope == "warp"):
if warp_id == 5:
tx_op(A_smem, A_smem, B_smem)
- elif exec_scope == "warpgroup":
+ elif T.constexpr(exec_scope == "warpgroup"):
if wg_id == 1:
tx_op(A_smem, A_smem, B_smem)
T.cuda.cta_sync()
@@ -309,7 +309,8 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
# in this test, use warp3/warpgroup1 to test
thr_str = 0 if exec_scope == "cta" else (128 if exec_scope == "warpgroup"
else 32 * 3)
a_shape = (n_threads, m, n)
- b_shape = (n_threads, m, n if rhs_kind == "region" else 1)
+ b_n = n if rhs_kind == "region" else 1
+ b_shape = (n_threads, m, b_n)
c_shape = a_shape
const = T.float16(1.25)
tx_op = {"add": Tx.add, "sub": Tx.sub, "mul": Tx.mul, "fdiv":
Tx.fdiv}[op_type]
@@ -329,7 +330,6 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
_bx = T.cta_id([1])
_tid = T.thread_id([256])
tid_in_scope = tid_in_scope_fn([n_threads])
- b_n = T.meta_var(n if rhs_kind == "region" else 1)
A_local = T.alloc_tensor((m, n), dtype, scope="local",
layout=TileLayout(S[(m, n)]))
C_local = T.alloc_tensor((m, n), dtype, scope="local",
layout=TileLayout(S[(m, n)]))
B_local = T.alloc_tensor((m, b_n), dtype, scope="local",
layout=TileLayout(S[(m, b_n)]))
@@ -338,25 +338,25 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
for i in T.serial(m):
for j in T.serial(n):
A_local[i, j] = A[tid_in_scope, i, j]
- if rhs_kind != "const":
+ if T.constexpr(rhs_kind != "const"):
for i in T.serial(m):
for j in T.serial(b_n):
B_local[i, j] = B[tid_in_scope, i, j]
- if exec_scope == "cta":
- if rhs_kind == "const":
+ if T.constexpr(exec_scope == "cta"):
+ if T.constexpr(rhs_kind == "const"):
tx_op(C_local, A_local, const)
else:
tx_op(C_local, A_local, B_local)
- elif exec_scope == "warpgroup":
+ elif T.constexpr(exec_scope == "warpgroup"):
if wg_id == 1:
- if rhs_kind == "const":
+ if T.constexpr(rhs_kind == "const"):
tx_op(C_local, A_local, const)
else:
tx_op(C_local, A_local, B_local)
else:
if warp_id == 3:
- if rhs_kind == "const":
+ if T.constexpr(rhs_kind == "const"):
tx_op(C_local, A_local, const)
else:
tx_op(C_local, A_local, B_local)
@@ -437,7 +437,7 @@ def test_binary_op_vectorized(input, storage_scope,
exec_scope, op_type, dtype):
T.device_entry()
_bx = T.cta_id([1])
tx = T.thread_id([thread_cnt])
- if storage_scope == "shared":
+ if T.constexpr(storage_scope == "shared"):
A_smem = T.alloc_tensor(a_shape, dtype, scope="shared",
layout=TileLayout(S[a_shape]))
B_smem = T.alloc_tensor(b_shape, dtype, scope="shared",
layout=TileLayout(S[b_shape]))
Tx.cta.copy(A_smem, A)
@@ -446,7 +446,7 @@ def test_binary_op_vectorized(input, storage_scope,
exec_scope, op_type, dtype):
tx_op(A_smem, A_smem, B_smem)
T.cuda.cta_sync()
Tx.cta.copy(A, A_smem)
- if storage_scope == "local":
+ if T.constexpr(storage_scope == "local"):
A_local = T.alloc_tensor(
a_shape[1:], dtype, scope="local",
layout=TileLayout(S[a_shape[1:]])
)
@@ -467,7 +467,7 @@ def test_binary_op_vectorized(input, storage_scope,
exec_scope, op_type, dtype):
T.device_entry()
_bx = T.cta_id([1])
tx = T.thread_id([thread_cnt])
- if storage_scope == "shared":
+ if T.constexpr(storage_scope == "shared"):
A_smem = T.alloc_tensor(a_shape, dtype, scope="shared",
layout=TileLayout(S[a_shape]))
B_smem = T.alloc_tensor(b_shape, dtype, scope="shared",
layout=TileLayout(S[b_shape]))
Tx.copy(A_smem, A)
@@ -476,7 +476,7 @@ def test_binary_op_vectorized(input, storage_scope,
exec_scope, op_type, dtype):
tx_op(A_smem, A_smem, B_smem)
T.cuda.cta_sync()
Tx.copy(A, A_smem)
- elif storage_scope == "local":
+ elif T.constexpr(storage_scope == "local"):
A_local = T.alloc_tensor(
a_shape[1:], dtype, scope="local",
layout=TileLayout(S[a_shape[1:]])
)
@@ -553,11 +553,11 @@ def test_binary_op_packed_f32x2_auto_dispatch(op_type):
)
Tx.copy(A_local, A[tx])
Tx.copy(B_local, B[tx])
- if op_type == "add":
+ if T.constexpr(op_type == "add"):
Tx.add(A_local, A_local, B_local)
- elif op_type == "sub":
+ elif T.constexpr(op_type == "sub"):
Tx.sub(A_local, A_local, B_local)
- elif op_type == "mul":
+ elif T.constexpr(op_type == "mul"):
Tx.mul(A_local, A_local, B_local)
Tx.copy(A[tx], A_local)
@@ -620,11 +620,11 @@ def test_binary_op_warpgroup_wg_local_layout(op_name):
lhs_row[i] = A[tid, i]
rhs_row[i] = B[tid, i]
out_row[i] = T.float32(0)
- if op_name == "add":
+ if T.constexpr(op_name == "add"):
Tx.wg.add(out, lhs, rhs)
- elif op_name == "sub":
+ elif T.constexpr(op_name == "sub"):
Tx.wg.sub(out, lhs, rhs)
- elif op_name == "mul":
+ elif T.constexpr(op_name == "mul"):
Tx.wg.mul(out, lhs, rhs)
out_row_1 = out.local(cols)
for i in T.serial(cols):
@@ -697,9 +697,9 @@ def
test_binary_op_warpgroup_wg_local_emits_packed_f32x2(op_name, ptx_op):
lhs_row[i] = A[tid, i]
rhs_row[i] = B[tid, i]
out_row[i] = T.float32(0)
- if op_name == "add":
+ if T.constexpr(op_name == "add"):
Tx.wg.add(out, lhs, rhs)
- elif op_name == "sub":
+ elif T.constexpr(op_name == "sub"):
Tx.wg.sub(out, lhs, rhs)
else:
Tx.wg.mul(out, lhs, rhs)
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 4f2e415d40..9849d812ca 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
@@ -82,9 +82,9 @@ def test_unary_op_shared(input, op_type, src_dtype,
dst_dtype):
A_smem = T.alloc_tensor(s_shape, src_dtype, scope="shared",
layout=s_layout)
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
T.cuda.cta_sync()
- if op_type == "zero":
+ if T.constexpr(op_type == "zero"):
Tx.cta.zero(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)])
- elif op_type == "sqrt":
+ elif T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt(A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)])
T.cuda.cta_sync()
Tx.cta.copy(A[tuple(copy_slice)], A_smem[tuple(copy_slice)])
@@ -104,9 +104,9 @@ def test_unary_op_shared(input, op_type, src_dtype,
dst_dtype):
B_smem = T.alloc_tensor(s_shape, dst_dtype, scope="shared",
layout=s_layout)
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
T.cuda.cta_sync()
- if op_type == "zero":
+ if T.constexpr(op_type == "zero"):
Tx.cta.zero(B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)])
- elif op_type == "sqrt":
+ elif T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt(B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)])
T.cuda.cta_sync()
Tx.cta.copy(B[tuple(map_slice_res)], B_smem[tuple(map_slice_res)])
@@ -169,10 +169,10 @@ def test_unary_op_shared_subcta_scope(exec_scope):
A_smem = T.alloc_tensor(g_shape, dtype, scope="shared",
layout=TileLayout(S[g_shape]))
Tx.cta.copy(A_smem, A)
T.cuda.cta_sync()
- if exec_scope == "warp":
+ if T.constexpr(exec_scope == "warp"):
if warp_id == 5:
Tx.warp.zero(A_smem, A_smem)
- elif exec_scope == "warpgroup":
+ elif T.constexpr(exec_scope == "warpgroup"):
if wg_id == 1:
Tx.wg.zero(A_smem, A_smem)
T.cuda.cta_sync()
@@ -261,30 +261,30 @@ def test_unary_op_shared_with_bias_scale(input, op_type,
bias_type, src_dtype, d
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
Tx.cta.copy(bias_smem[tuple(copy_slice)], bias[tuple(copy_slice)])
T.cuda.cta_sync()
- if bias_type == "const":
- if op_type == "sqrt":
+ if T.constexpr(bias_type == "const"):
+ if T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=const_bias,
)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.cta.exp_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=const_bias,
)
- elif bias_type == "region":
- if op_type == "sqrt":
+ elif T.constexpr(bias_type == "region"):
+ if T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=bias_smem[tuple(map_slice_a)],
)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.cta.exp_with_scale_bias(
A_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
@@ -310,30 +310,30 @@ def test_unary_op_shared_with_bias_scale(input, op_type,
bias_type, src_dtype, d
Tx.cta.copy(A_smem[tuple(copy_slice)], A[tuple(copy_slice)])
Tx.cta.copy(bias_smem[tuple(copy_slice)], bias[tuple(copy_slice)])
T.cuda.cta_sync()
- if bias_type == "const":
- if op_type == "sqrt":
+ if T.constexpr(bias_type == "const"):
+ if T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=const_bias,
)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.cta.exp_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=const_bias,
)
- elif bias_type == "region":
- if op_type == "sqrt":
+ elif T.constexpr(bias_type == "region"):
+ if T.constexpr(op_type == "sqrt"):
Tx.cta.sqrt_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
scale=scale,
bias=bias_smem[tuple(map_slice_a)],
)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.cta.exp_with_scale_bias(
B_smem[tuple(map_slice_res)],
A_smem[tuple(map_slice_a)],
@@ -502,13 +502,13 @@ def test_unary_op_local(input, op_type, src_dtype,
dst_dtype):
# unary op
acc_view = acc.view(*acc_shape, layout=acc_layout)
res_view = res.view(*red_shape, layout=acc_layout)
- if op_type == "reciprocal":
+ if T.constexpr(op_type == "reciprocal"):
Tx.warp.reciprocal(res_view, acc_view)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.warp.exp(res_view, acc_view)
- elif op_type == "exp2":
+ elif T.constexpr(op_type == "exp2"):
Tx.warp.exp2(res_view, acc_view)
- elif op_type == "log2":
+ elif T.constexpr(op_type == "log2"):
Tx.warp.log2(res_view, acc_view)
# write res into B
@@ -655,15 +655,15 @@ def test_unary_op_local_with_bias_scale(input, op_type,
bias_type, src_dtype, ds
acc_view = acc.view(*acc_shape, layout=acc_layout)
res_view = res.view(*red_shape, layout=acc_layout)
bias_view = bias_local.view(*bias_shape, layout=acc_layout)
- if bias_type == "const":
- if op_type == "sqrt":
+ if T.constexpr(bias_type == "const"):
+ if T.constexpr(op_type == "sqrt"):
Tx.warp.sqrt_with_scale_bias(res_view, acc_view, scale=scale,
bias=const_bias)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.warp.exp_with_scale_bias(res_view, acc_view, scale=scale,
bias=const_bias)
- elif bias_type == "region":
- if op_type == "sqrt":
+ elif T.constexpr(bias_type == "region"):
+ if T.constexpr(op_type == "sqrt"):
Tx.warp.sqrt_with_scale_bias(res_view, acc_view, scale=scale,
bias=bias_view)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.warp.exp_with_scale_bias(res_view, acc_view, scale=scale,
bias=bias_view)
# write res into B
@@ -733,13 +733,13 @@ def test_unary_op_vectorized(shape, op_type, exec_scope,
storage_scope):
T.device_entry()
_bx = T.cta_id([1])
tx = T.thread_id([128])
- if storage_scope == "shared":
+ if T.constexpr(storage_scope == "shared"):
a_smem = T.alloc_tensor(
shape, dtype=dtype, layout=TileLayout(S[shape]), scope="shared"
)
Tx.fill(a_smem[tx], value)
Tx.copy(A[tx], a_smem[tx])
- elif storage_scope == "local":
+ elif T.constexpr(storage_scope == "local"):
a_local = T.alloc_tensor(
shape[1:], dtype=dtype, layout=TileLayout(S[shape[1:]]),
scope="local"
)
@@ -752,7 +752,7 @@ def test_unary_op_vectorized(shape, op_type, exec_scope,
storage_scope):
T.device_entry()
_bx = T.cta_id([1])
_tid = T.thread_id([128])
- if storage_scope == "shared":
+ if T.constexpr(storage_scope == "shared"):
a_smem = T.alloc_tensor(
shape, dtype=dtype, layout=TileLayout(S[shape]), scope="shared"
)
@@ -794,15 +794,15 @@ def test_unary_op_local_thread_wise(op_type, dtype):
local_shape, dtype, scope="local",
layout=TileLayout(S[local_shape])
)
Tx.copy(a_local, A[tid])
- if op_type == "zero":
+ if T.constexpr(op_type == "zero"):
Tx.zero(a_local, a_local)
- elif op_type == "sqrt":
+ elif T.constexpr(op_type == "sqrt"):
Tx.sqrt(a_local, a_local)
- elif op_type == "reciprocal":
+ elif T.constexpr(op_type == "reciprocal"):
Tx.reciprocal(a_local, a_local)
- elif op_type == "exp":
+ elif T.constexpr(op_type == "exp"):
Tx.exp(a_local, a_local)
- elif op_type == "silu":
+ elif T.constexpr(op_type == "silu"):
Tx.silu(a_local, a_local)
Tx.copy(A[tid], a_local)
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py
b/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py
index 214f2cc354..ce68054fbb 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/gemm/test_gemm_mma_m16n8k_.py
@@ -33,6 +33,7 @@ is guarded by ``requires_cuda`` since it needs a real device.
import numpy as np
import pytest
+from tvm_ffi import structural_walk
import tvm
import tvm.testing
@@ -278,7 +279,7 @@ def _build_tiled_numeric(Mt, Nt, Kt, kinst, beta, dtype):
kt * kinst + kHi * 8 + 2 * (lane % 4) + kp,
nt * 8 + lane // 4,
]
- if beta == 1.0:
+ if T.constexpr(beta == 1.0):
C_reg = C_f.local(Mt, Nt, 2, 2)
for mt, nt, rM, rN in T.grid(Mt, Nt, 2, 2):
C_reg[mt, nt, rM, rN] = C_g[
@@ -318,7 +319,7 @@ def _build_transpose_numeric(transpose_A, transpose_B,
dtype="float16"):
B_f = T.alloc_tensor(B_shape, dtype, scope="local", layout=Bl)
D_f = T.alloc_tensor((16, 8), "float32", scope="local", layout=D_FRAG)
A_reg = A_f.local(2, 2, 2)
- if transpose_A:
+ if T.constexpr(transpose_A):
# A_KM_FRAG: buffer is [K, M].
for kHi, rM, kp in T.grid(2, 2, 2):
A_reg[kHi, rM, kp] = A_g[2 * (lane % 4) + kp + 8 * kHi, lane
// 4 + 8 * rM]
@@ -327,7 +328,7 @@ def _build_transpose_numeric(transpose_A, transpose_B,
dtype="float16"):
for kHi, rM, kp in T.grid(2, 2, 2):
A_reg[kHi, rM, kp] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) +
kp + 8 * kHi]
B_reg = B_f.local(2, 2)
- if transpose_B:
+ if T.constexpr(transpose_B):
# B_NK_FRAG buffer is [N, K].
for kHi, kp in T.grid(2, 2):
B_reg[kHi, kp] = B_g[lane // 4, 2 * (lane % 4) + kp + 8 * kHi]
@@ -372,35 +373,51 @@ def test_cuda_gemm_mma_variant_is_registered():
def test_cuda_gemm_mma_lowers_to_mma_sync(dtype):
"""beta=0: the dispatch clears D, then issues a single accumulating mma
with
the registers laid out in the fixed PTX fragment order."""
- script = _lower(_build_gemm(alpha=1.0, beta=0.0,
dtype=dtype))["main"].script()
-
- assert "T.ptx.mma(" in script
- assert "m16n8k16" in script
+ func = _lower(_build_gemm(alpha=1.0, beta=0.0, dtype=dtype))["main"]
+ calls, stores = [], []
+ structural_walk(func.body, [(tvm.ir.Call, calls.append),
(tvm.tirx.TensorStore, stores.append)])
+ mma_calls = [call for call in calls if call.op.name == "tirx.ptx.mma"]
+ assert len(mma_calls) == 1
+ operands = mma_calls[0].args
+ assert operands[16].value == "m16n8k16"
# beta == 0 clears the accumulator before the K loop.
- assert "T.float32(0" in script
+ assert len(stores) == 1
+ assert isinstance(stores[0].value, tvm.tirx.FloatImm) and
stores[0].value.value == 0
+ assert stores[0].buffer.same_as(operands[0].source)
# D accumulator: c_id = 2*rM + rN -> regs 0..3.
- for r in range(4):
- assert f"d_local[{r}]" in script
# A and B fragments are packed two elements per b32, so the instruction
# indexes a uint32 view: the element strides above halve into word strides.
# A: b32 = rM + 2*kHi (kHi outer) -> words 0..3.
- for r in range(4):
- assert f"a_words[{r}]" in script
# B: b32 = kHi -> words 0, 1.
- for r in (0, 1):
- assert f"b_words[{r}]" in script
+ for fragment, dtype in [
+ (operands[:4], "float32"),
+ (operands[4:8], "uint32"),
+ (operands[8:10], "uint32"),
+ ]:
+ for r, load in enumerate(fragment):
+ assert isinstance(load, tvm.ir.TensorLoad)
+ assert load.source.same_as(fragment[0].source)
+ assert load.ty.dtype == dtype
+ assert len(load.indices) == 1 and int(load.indices[0]) == r
+ for dst, src in zip(operands[:4], operands[10:14]):
+ tvm.ir.assert_structural_equal(dst, src)
@pytest.mark.gpu
def test_cuda_gemm_mma_accumulates_c_when_beta_one():
"""beta=1: the accumulator is initialized by copying C instead of
zeroing."""
- script = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"].script()
-
- assert "T.ptx.mma(" in script
- assert "m16n8k16" in script
+ func = _lower(_build_gemm(alpha=1.0, beta=1.0))["main"]
+ calls, stores = [], []
+ structural_walk(func.body, [(tvm.ir.Call, calls.append),
(tvm.tirx.TensorStore, stores.append)])
+ mma_calls = [call for call in calls if call.op.name == "tirx.ptx.mma"]
+ assert len(mma_calls) == 1
# The init reads C into D; nothing is zeroed.
- assert "c_local[" in script
- assert "T.float32(0" not in script
+ assert len(stores) == 1
+ init = stores[0]
+ assert init.buffer.same_as(mma_calls[0].args[0].source)
+ assert isinstance(init.value, tvm.ir.TensorLoad)
+ assert not init.buffer.same_as(init.value.source)
+ tvm.ir.assert_structural_equal(init.indices, init.value.indices)
def test_cuda_gemm_mma_rejects_nonunit_alpha():
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
index 7c1540c266..7581d3303f 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
@@ -2978,7 +2978,7 @@ def _build_smem_desc_kernel(smem_desc,
weight_stationary=False, pass_descI=False
T.cuda.mbarrier_wait(tma_mbar.ptr_to([0]), 0)
T.cuda.cta_sync()
if tid_in_wg == 0:
- if pass_descI:
+ if T.constexpr(pass_descI):
desc_i: T.uint32
T.cuda.tcgen05.encode_instr_descriptor(
T.address_of(desc_i), # noqa: F821
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py
b/tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py
index 3da84058b6..2a18e42448 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/reduction/test_reduction.py
@@ -81,24 +81,24 @@ def test_reduction_shared(
B_smem = T.alloc_tensor(s_shape_dst, dtype, scope="shared",
layout=s_layout_dst)
Tx.cta.copy(A_smem[tuple(copy_slice_src)], A[tuple(copy_slice_src)])
- if accum:
+ if T.constexpr(accum):
Tx.cta.copy(B_smem[tuple(copy_slice_dst)],
B[tuple(copy_slice_dst)])
T.cuda.cta_sync()
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.cta.sum(
B_smem[tuple(reduce_slice_dst)],
A_smem[tuple(reduce_slice_src)],
axes=axes,
accum=accum,
)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.cta.max(
B_smem[tuple(reduce_slice_dst)],
A_smem[tuple(reduce_slice_src)],
axes=axes,
accum=accum,
)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.cta.min(
B_smem[tuple(reduce_slice_dst)],
A_smem[tuple(reduce_slice_src)],
@@ -181,15 +181,15 @@ def test_reduction_shared_subscope(exec_scope, op_type,
accum):
A_smem = T.alloc_tensor(list(src_shape), dtype, scope="shared",
layout=s_layout_src)
B_smem = T.alloc_tensor(list(dst_shape), dtype, scope="shared",
layout=s_layout_dst)
Tx.cta.copy(A_smem, A)
- if accum:
+ if T.constexpr(accum):
Tx.cta.copy(B_smem, B)
T.cuda.cta_sync()
if warp_id == 5:
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.warp.min(B_smem, A_smem, axes=axes, accum=accum)
T.cuda.cta_sync()
Tx.cta.copy(B, B_smem)
@@ -208,15 +208,15 @@ def test_reduction_shared_subscope(exec_scope, op_type,
accum):
A_smem = T.alloc_tensor(list(src_shape), dtype, scope="shared",
layout=s_layout_src)
B_smem = T.alloc_tensor(list(dst_shape), dtype, scope="shared",
layout=s_layout_dst)
Tx.cta.copy(A_smem, A)
- if accum:
+ if T.constexpr(accum):
Tx.cta.copy(B_smem, B)
T.cuda.cta_sync()
if wg_id == 0:
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.wg.sum(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.wg.max(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.wg.min(B_smem, A_smem, axes=axes, accum=accum)
T.cuda.cta_sync()
Tx.cta.copy(B, B_smem)
@@ -234,15 +234,15 @@ def test_reduction_shared_subscope(exec_scope, op_type,
accum):
A_smem = T.alloc_tensor(list(src_shape), dtype, scope="shared",
layout=s_layout_src)
B_smem = T.alloc_tensor(list(dst_shape), dtype, scope="shared",
layout=s_layout_dst)
Tx.cta.copy(A_smem, A)
- if accum:
+ if T.constexpr(accum):
Tx.cta.copy(B_smem, B)
T.cuda.cta_sync()
if _tid == 65:
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.sum(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.max(B_smem, A_smem, axes=axes, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.min(B_smem, A_smem, axes=axes, accum=accum)
T.cuda.cta_sync()
Tx.cta.copy(B, B_smem)
@@ -339,16 +339,16 @@ def test_reduction_local_thread_wise(src_shape,
dst_shape, axes, op_type, accum)
idx = T.meta_var(decompose_flat(i, src_shape))
A_local[tuple(idx)] = A[tuple(idx)]
- if accum:
+ if T.constexpr(accum):
for i in T.serial(dst_total):
idx = T.meta_var(decompose_flat(i, dst_shape))
B_local[tuple(idx)] = B[tuple(idx)]
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.sum(B_local, A_local, axes=axes, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.max(B_local, A_local, axes=axes, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.min(B_local, A_local, axes=axes, accum=accum)
for i in T.serial(dst_total):
@@ -468,31 +468,31 @@ def test_reduction_local_view_basic(inner_dims, dst_dims,
axes, accum, slice_end
for i in T.serial(src_local_total):
idx = T.meta_var(decompose_flat(i, inner_dims))
acc[(0, *list(idx))] = A[(lane_id, *list(idx))]
- if accum:
+ if T.constexpr(accum):
for i in T.serial(dst_local_total):
idx = T.meta_var(decompose_flat(i, dst_dims))
red[(0, *list(idx))] = B[(lane_id, *list(idx))]
acc_view = acc.view(*src_shape, layout=acc_view_layout)
red_view = red.view(*dst_shape, layout=red_view_layout)
- if slice_end is not None:
- if op_type == "sum":
+ if T.constexpr(slice_end is not None):
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(
red_view, acc_view[:, slice_end // 2 : slice_end],
axes=axes, accum=accum
)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(
red_view, acc_view[:, slice_end // 2 : slice_end],
axes=axes, accum=accum
)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.warp.min(
red_view, acc_view[:, slice_end // 2 : slice_end],
axes=axes, accum=accum
)
else:
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(red_view, acc_view, axes=axes, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(red_view, acc_view, axes=axes, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.warp.min(red_view, acc_view, axes=axes, accum=accum)
for i in T.serial(dst_local_total):
idx = T.meta_var(decompose_flat(i, dst_dims))
@@ -599,7 +599,7 @@ def test_reduction_local_view_complex(n_groups, n_warps,
op_type, dtype, shuffle
]
# Pre-load B into red for accumulation
- if accum:
+ if T.constexpr(accum):
for i in T.unroll(2):
red[i] = B[
wg_id * 64 + warp_id_in_wg * 16 + i * 8 + lane_id // 4,
@@ -607,19 +607,19 @@ def test_reduction_local_view_complex(n_groups, n_warps,
op_type, dtype, shuffle
]
acc_view = acc.view(*acc_shape, layout=acc_layout)
red_view = red.view(*red_shape, layout=red_layout)
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(red_view, acc_view, thread_reduce=shuffle, accum=accum)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(red_view, acc_view, thread_reduce=shuffle, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.warp.min(red_view, acc_view, thread_reduce=shuffle, accum=accum)
# perform an additional shuffle step if not shuffled above
- if not shuffle:
- if op_type == "sum":
+ if T.constexpr(not shuffle):
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(red_view, red_view, thread_reduce=True)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(red_view, red_view, thread_reduce=True)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.warp.min(red_view, red_view, thread_reduce=True)
# Write red into B
for i in T.unroll(2):
@@ -698,13 +698,13 @@ def
test_reduction_local_optimized_3input_maxmin(reduction_len, op_type, accum):
A_local[i] = A[i]
# Initialize B_local for accum test
- if accum:
+ if T.constexpr(accum):
B_local[0] = B[0]
# Thread-level reduction
- if op_type == "max":
+ if T.constexpr(op_type == "max"):
Tx.max(B_local, A_local, accum=accum)
- elif op_type == "min":
+ elif T.constexpr(op_type == "min"):
Tx.min(B_local, A_local, accum=accum)
# Store result to global
@@ -771,7 +771,7 @@ def
test_reduction_local_optimized_packed_add_sum(reduction_len, accum):
A_local[i] = A[i]
# Initialize B_local for accum test
- if accum:
+ if T.constexpr(accum):
B_local[0] = B[0]
# Thread-level sum reduction
@@ -845,9 +845,9 @@ def test_reduction_op_warp_shuffle(op_type, dtype):
src_local[0] = A[lane_id]
src_view = src_local.view(N, layout=src_layout)
dst_view = dst_local.view(1, layout=dst_layout)
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(dst_view, src_view)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(dst_view, src_view)
B[lane_id] = dst_local[0]
# fmt: on
@@ -917,9 +917,9 @@ def test_reduction_op_warp_shuffle_multi_elem(op_type,
dtype):
src_local[i] = A[lane_id * ELEMS_PER_THREAD + i]
src_view = src_local.view(TOTAL, layout=src_layout)
dst_view = dst_local.view(ELEMS_PER_THREAD, layout=dst_layout)
- if op_type == "sum":
+ if T.constexpr(op_type == "sum"):
Tx.warp.sum(dst_view, src_view)
- elif op_type == "max":
+ elif T.constexpr(op_type == "max"):
Tx.warp.max(dst_view, src_view)
for i in T.serial(ELEMS_PER_THREAD):
B[i] = dst_local[i]
@@ -1115,7 +1115,7 @@ def test_reduction_warpgroup_wg_local_layout(op_name):
src_local = src.local(cols)
for i in T.serial(cols):
src_local[i] = A[tid, i]
- if op_name == "sum":
+ if T.constexpr(op_name == "sum"):
Tx.wg.sum(dst, src, axes=[-1], accum=False)
else:
Tx.wg.max(dst, src, axes=[-1], accum=False)
diff --git a/tests/python/tirx/test_op_namespace_cleanup.py
b/tests/python/tirx/test_op_namespace_cleanup.py
index f5ffce8333..33cc63a5fd 100644
--- a/tests/python/tirx/test_op_namespace_cleanup.py
+++ b/tests/python/tirx/test_op_namespace_cleanup.py
@@ -105,6 +105,36 @@ def test_builtin_expression_ops_are_not_tile_primitives():
assert fma.op.name == "tirx.fma"
+def test_tensormap_encoding_accepts_attribute_keywords_and_defaults():
+ @T.function
+ def encode(descriptor: T.TensorMap, data: T.handle):
+ T.tensormap_encode_tiled(
+ descriptor,
+ data,
+ 64,
+ 32,
+ 128,
+ 64,
+ 32,
+ 1,
+ 1,
+ descriptor_dtype="bfloat16",
+ rank=2,
+ swizzle=3,
+ )
+
+ (call,) = [c for c in _expr_calls(encode) if c.op.name ==
"tirx.tensormap_encode_tiled"]
+ assert call.op.attrs_type_key == "tirx.TensorMapEncodeTiledAttr"
+ assert str(call.attrs.descriptor_dtype) == "bfloat16"
+ assert call.attrs.rank == 2
+ assert call.attrs.swizzle == 3
+ assert call.attrs.interleave == call.attrs.l2_promotion ==
call.attrs.oob_fill == 0
+ assert call.attrs.force_cu_dtype == -1
+ call.validate()
+ reparsed = tvm.script.from_source(encode.script(), extra_vars={"T": T})
+ assert_structural_equal(encode, reparsed)
+
+
def test_kernel_replace_point_is_builtin_marker_not_tile_primitive():
assert _op_attr("tirx.tvm_kernel_replace_point", "TIRxOpCategory") ==
"builtin"
assert "tirx.tile.tvm_kernel_replace_point" not in Op.list_op_names()
diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py
b/tests/python/tirx/transform/test_transform_lower_tirx.py
index e0c60d951e..5d54c99ae2 100644
--- a/tests/python/tirx/transform/test_transform_lower_tirx.py
+++ b/tests/python/tirx/transform/test_transform_lower_tirx.py
@@ -1439,6 +1439,36 @@ def
test_lower_exec_context_tracks_cta_id_in_pair_after_axis_predicate():
assert _int_triple(seen[0], "cby") == (1, 1, 2)
[email protected]("offset", [0, 16])
+def test_lower_remaps_tensor_memory_address_metadata(offset):
+ @T.function(private=True)
+ def before():
+ T.device_entry()
+ addr = T.alloc_shared((1,), "uint32")
+ for i in T.serial(2):
+ tensor = T.decl_tensor(
+ (64,),
+ "float32",
+ scope="tmem",
+ allocated_addr=addr[0] + T.uint32(offset) + T.cast(i,
"uint32"),
+ )
+ T.evaluate(tensor.data)
+
+ with tvm.target.Target("cuda"):
+ lowered = LowerTIRx()(tvm.IRModule({"main": before}))["main"]
+ assert not tvm.tirx.analysis.undefined_vars(lowered.body, lowered.params)
+ tvm.tirx.analysis.verify_well_formed(lowered)
+ bindings = []
+ tvm_ffi.structural_walk(lowered.body, (tvm.tirx.Bind, bindings.append))
+ address = next(node.var for node in bindings if node.var.scope() ==
"shared")
+ tensor = next(node for node in bindings if node.var.scope() == "tmem")
+ loads = []
+ tvm_ffi.structural_walk(tensor.var.ty.allocated_addr, (tvm.ir.TensorLoad,
loads.append))
+ assert len(loads) == 1
+ assert loads[0].source.same_as(address)
+ tvm.ir.assert_structural_equal(tensor.var.ty, tensor.value.ty)
+
+
def test_lower_buffer_offset():
@T.function(private=True)
def before():
