[Fix] Harden row-expand layouts and clean up legacy tooling 🤖

Fold contiguous leading row-expand dimensions into rows while validating the 256-byte row, packed-scalar, workspace, dtype, and contiguity contracts.

Remove the obsolete string-based init_flag/clear_flag frontend and codegen path, and make set_env.sh robust when sourced from Bash or Zsh without duplicating PYTHONPATH.

Keep CI coverage focused with shared-helper validation, one two-backend multistage numerical regression, and a legacy-export check.

Author: platelett

set_env.sh

@@ -1,8 +1,20 @@
 #!/bin/bash
 
-TL_ROOT=$(readlink -f "${BASH_SOURCE[0]}")
-export TL_ROOT=$(dirname "$TL_ROOT")
-export PYTHONPATH=${TL_ROOT}:$PYTHONPATH
+if [ -n "${ZSH_VERSION:-}" ]; then
+    tilelang_env_script_path="${(%):-%N}"
+else
+    tilelang_env_script_path="${BASH_SOURCE[0]}"
+fi
+
+TL_ROOT=$(dirname "$(readlink -f "$tilelang_env_script_path")")
+export TL_ROOT
+case "${PYTHONPATH:-}" in
+    "$TL_ROOT"|"$TL_ROOT":*) ;;
+    *) PYTHONPATH="$TL_ROOT${PYTHONPATH:+:$PYTHONPATH}" ;;
+esac
+export PYTHONPATH
 
 # disable the import of tvm when using torch_npu
 export ACL_OP_INIT_MODE=1
+
+unset tilelang_env_script_path

src/target/codegen_ascend.cc

@@ -770,20 +770,6 @@ void CodeGenTileLangAscend::VisitStmt_(const AttrStmtNode *op) {
     }
     this->VisitStmt(op->body);
     return;
-  } else if (op->attr_key == "init_flag" || op->attr_key == "clear_flag") {
-    const StringImmNode *instn = op->value.as<StringImmNode>();
-
-    std::string inst = std::string(instn->value);
-    size_t st = 0;
-    for (size_t i = 0; i < inst.size(); ++i) {
-      if (inst[i] == '\n') {
-        this->PrintIndent();
-        stream << inst.substr(st, i - st) << "\n";
-        st = i + 1;
-      }
-    }
-    this->VisitStmt(op->body);
-    return;
   } else if (op->attr_key == "resource_scope") {
     auto resource_id = Downcast<IntImm>(op->value)->value;
     auto resource_name = resource_id == 0 ? "AIC" : "AIV";

testing/python/language/test_tilelang_ascend_language_elementwise.py

@@ -5137,6 +5137,121 @@ def row_expand_div_experiment_kernel(M, N, dtype="float16"):
     return _row_expand_binop_experiment_kernel(M, N, "row_expand_div_experiment", dtype)
 
 
+def row_expand_mul_multistage_region_kernel():
+    stages = 2
+    rows_per_stage = 8
+    cols = 64
+    lanes = 8
+
+    @T.prim_func
+    def main(
+        A: T.Tensor((stages, rows_per_stage, cols), "float"),  # type: ignore
+        S: T.Tensor((stages, rows_per_stage, lanes), "float"),  # type: ignore
+        C: T.Tensor((stages, rows_per_stage, cols), "float"),  # type: ignore
+    ):
+        with T.Kernel(1, is_npu=True):
+            a_ring = T.alloc_ub((stages + 1, rows_per_stage, cols), "float")
+            s_ring = T.alloc_ub((stages + 1, rows_per_stage, lanes), "float")
+            c_ring = T.alloc_ub((stages + 1, rows_per_stage, cols), "float")
+
+            for stage in T.serial(stages):
+                T.copy(A[stage, :, :], a_ring[stage + 1, :, :])
+                T.copy(S[stage, :, :], s_ring[stage + 1, :, :])
+            T.tile.row_expand_mul_experiment(
+                c_ring[1:3, :, :],
+                a_ring[1:3, :, :],
+                s_ring[1:3, :, :],
+            )
+            for stage in T.serial(stages):
+                T.copy(c_ring[stage + 1, :, :], C[stage, :, :])
+
+    return main
+
+
+def test_row_expand_experiment_folds_contiguous_leading_dimensions():
+    op = T.tile.row_expand_mul_experiment
+    stage = tir.Var("stage", "int32")
+    dst_ring = tir.decl_buffer((3, 8, 64), "float32", scope="shared.ub")
+    src_ring = tir.decl_buffer((3, 8, 64), "float32", scope="shared.ub")
+    scalar_ring = tir.decl_buffer((3, 8, 8), "float32", scope="shared.ub")
+
+    singleton = op(dst_ring[stage, :, :], src_ring[stage, :, :], scalar_ring[stage, :, :])
+    multistage = op(dst_ring[1:3, :, :], src_ring[1:3, :, :], scalar_ring[1:3, :, :])
+    whole_buffer = op(dst_ring, src_ring, scalar_ring)
+
+    assert int(singleton.args[1].args[3]) == 8 * 64
+    assert int(singleton.args[3].args[3]) == 8 * 8
+    assert int(multistage.args[1].args[3]) == 2 * 8 * 64
+    assert int(multistage.args[3].args[3]) == 2 * 8 * 8
+    assert int(whole_buffer.args[1].args[3]) == 3 * 8 * 64
+    assert int(whole_buffer.args[3].args[3]) == 3 * 8 * 8
+
+
+def test_row_expand_experiment_rejects_invalid_folded_regions():
+    op = T.tile.row_expand_mul_experiment
+    dst = tir.decl_buffer((64, 64), "float32", scope="shared.ub")
+    src = tir.decl_buffer((64, 64), "float32", scope="shared.ub")
+    extra_rows = tir.decl_buffer((2, 64, 8), "float32", scope="shared.ub")
+
+    mismatch = r"src1 scalar count must match dst rows: src1=128, dst\[0\]=64"
+    with pytest.raises(ValueError, match=mismatch):
+        op(dst, src, extra_rows)
+
+    scalar_ring = tir.decl_buffer((2, 128, 8), "float32", scope="shared.ub")
+    with pytest.raises(ValueError, match="outer dimensions must be contiguous"):
+        op(dst, src, scalar_ring[0:2, 0:64, :])
+
+    with pytest.raises(ValueError, match="requires a 256-byte last dimension"):
+        op(
+            tir.decl_buffer((8, 128), "float32", scope="shared.ub"),
+            tir.decl_buffer((8, 128), "float32", scope="shared.ub"),
+            tir.decl_buffer((8, 8), "float32", scope="shared.ub"),
+        )
+
+
+def test_row_expand_experiment_validates_src1_and_tmp_layouts():
+    op = T.tile.row_expand_mul_experiment
+    dst = tir.decl_buffer((16, 64), "float32", scope="shared.ub")
+    src = tir.decl_buffer((16, 64), "float32", scope="shared.ub")
+    packed = tir.decl_buffer((16, 8), "float32", scope="shared.ub")
+    scalar = tir.decl_buffer((16,), "float32", scope="shared.ub")
+    tmp = tir.decl_buffer((16, 8), "float32", scope="shared.ub")
+
+    op(dst, src, packed)
+    op(dst, src, scalar, tmp)
+
+    with pytest.raises(ValueError, match="packed src1 shape"):
+        op(dst, src, scalar)
+    with pytest.raises(ValueError, match="scalar src1 shape"):
+        op(dst, src, packed, tmp)
+    with pytest.raises(ValueError, match="tmp must contain 128 elements"):
+        op(dst, src, scalar, tir.decl_buffer((15, 8), "float32", scope="shared.ub"))
+
+    strided_packed = tir.decl_buffer((16, 16), "float32", scope="shared.ub")
+    with pytest.raises(ValueError, match="must be contiguous when flattened"):
+        op(dst, src, strided_packed[:, :8])
+
+
+@pytest.mark.parametrize("target", ["ascendc", "pto"])
+def test_row_expand_mul_experiment_multistage_region(target):
+    func = tilelang.compile(
+        row_expand_mul_multistage_region_kernel(),
+        out_idx=[-1],
+        pass_configs=pass_configs,
+        target=target,
+    )
+
+    a = torch.arange(1, 2 * 8 * 64 + 1, dtype=torch.float32).reshape(2, 8, 64) / 64
+    scalars = torch.arange(1, 2 * 8 + 1, dtype=torch.float32).reshape(2, 8) / 4
+    packed_scalars = scalars.unsqueeze(-1).expand(2, 8, 8).contiguous()
+
+    c = func(a.npu(), packed_scalars.npu())
+    torch.npu.synchronize()
+    ref_c = a * scalars.unsqueeze(-1)
+
+    torch.testing.assert_close(c.cpu(), ref_c, rtol=1e-5, atol=1e-5)
+
+
 @pytest.mark.parametrize("dtype,shape", [("float16", (16, 128)), ("float", (16, 64))])
 @pytest.mark.parametrize("target", ["ascendc", "pto"])
 def test_row_expand_mul_experiment(dtype, target, shape):

testing/python/language/test_tilelang_ascend_language_legacy_flags.py

@@ -0,0 +1,8 @@
+"""Regression tests for removal of legacy string-based flag helpers."""
+
+import tilelang.language as T
+
+
+def test_legacy_flag_helpers_are_not_exported():
+    assert not hasattr(T, "init_flag")
+    assert not hasattr(T, "clear_flag")

tilelang/language/__init__.py

@@ -173,32 +173,6 @@ def import_source(source: str | None = None):
     return block_attr({"pragma_import_c": source}) if source is not None else None
 
 
-def init_flag(fmap):
-    inst = ""
-    for src, d in fmap.items():
-        for dst, stages in d.items():
-            for stage in stages:
-                inst += f"AscendC::SetFlag<AscendC::HardEvent::{src}_{dst}>({stage});\n"
-
-    return attr(None, "init_flag", inst)
-
-
-def clear_flag(fmap):
-    inst = ""
-    for src, d in fmap.items():
-        for dst, stages in d.items():
-            for stage in stages:
-                inst += f"AscendC::WaitFlag<AscendC::HardEvent::{src}_{dst}>({stage});\n"
-
-    @macro
-    def _get_inst():
-        with attr(None, "clear_flag", inst):
-            call_extern("handle", "...")
-
-    # return attr(call_extern("handle", "..."), "clear_flag", inst)
-    return _get_inst()
-
-
 def npu_use_swizzle(cid, m, n, k, block_m, block_n, off=1, dir=0, in_loop=False):
     # If order is row, use rasterization2DRow, otherwise use rasterization2DColumn
     # The panel size is the number of threads in a warp

tilelang/language/ascend_tile.py

@@ -76,6 +76,7 @@ def _handle_buffer_region_2d(br: BufferRegion, mask):
 
     Leading dimensions are folded into rows; the innermost dimension is kept as cols.
     """
+    _validate_buffer_region_outer_contiguity(br)
     bf = br.buffer
     indices = [x.min for x in br.region]
     offset = bf.offset_of(indices)[0]
@@ -2300,6 +2301,49 @@ def _shapes_equal(shape1, shape2) -> bool:
     return all(_const_equal(x, y) for x, y in zip(shape1, shape2))
 
 
+def _fold_nd_shape_to_2d(shape):
+    """Fold all leading dimensions of a shape into a single row dimension."""
+    shape = list(shape)
+    if len(shape) <= 1:
+        return shape
+    return [math.prod(shape[:-1]), shape[-1]]
+
+
+def _validate_buffer_region_outer_contiguity(br: BufferRegion) -> None:
+    """Validate that folding a region's outer dimensions does not cross gaps.
+
+    The final dimension is a row window and may be narrower than the physical
+    buffer because row-expand codegen retains the physical row stride. Once an
+    outer dimension spans multiple entries, however, every following outer
+    dimension must be selected in full.
+    """
+    spans_multiple_entries = False
+    for axis, region in enumerate(br.region[:-1]):
+        if spans_multiple_entries:
+            starts_at_zero = _const_equal(region.min, 0)
+            has_full_extent = _const_equal(region.extent, br.buffer.shape[axis])
+            full_axis = starts_at_zero and has_full_extent
+            if not full_axis:
+                requirement = "BufferRegion outer dimensions must be contiguous"
+                raise ValueError(f"{requirement}; axis {axis} is not selected in full.")
+        elif not _is_const_one(region.extent):
+            spans_multiple_entries = True
+
+
+def _validate_buffer_region_flat_contiguity(br: BufferRegion) -> None:
+    """Validate that a region can be consumed as one flat contiguous stream."""
+    spans_multiple_entries = False
+    for axis, region in enumerate(br.region):
+        if spans_multiple_entries:
+            starts_at_zero = _const_equal(region.min, 0)
+            has_full_extent = _const_equal(region.extent, br.buffer.shape[axis])
+            if not (starts_at_zero and has_full_extent):
+                requirement = "BufferRegion must be contiguous when flattened"
+                raise ValueError(f"{requirement}; axis {axis} is not selected in full.")
+        elif not _is_const_one(region.extent):
+            spans_multiple_entries = True
+
+
 def _row_expand_binop_experiment(
     dst,
     src0,
@@ -2319,31 +2363,54 @@ def _row_expand_binop_experiment(
         dst_ptr, dst_shape = _handle_buffer_region_2d(dst, "w")
     else:
         dst_ptr = dst.access_ptr("w")
-        dst_shape = list(dst.shape[-2:])
+        dst_shape = _fold_nd_shape_to_2d(dst.shape)
 
     if isinstance(src0, BufferRegion):
         src0_ptr, src0_shape = _handle_buffer_region_2d(src0, "r")
     else:
         src0_ptr = src0.access_ptr("r")
-        src0_shape = list(src0.shape[-2:])
+        src0_shape = _fold_nd_shape_to_2d(src0.shape)
 
     if isinstance(src1, BufferRegion):
-        src1_ptr, src1_nd_extent = _handle_buffer_region(src1, "r")
-        src1_full_shape = [src1_nd_extent[-1]] if len(src1_nd_extent) >= 2 else src1_nd_extent
+        src1_ptr, src1_full_shape = _handle_buffer_region_2d(src1, "r")
     else:
         src1_ptr = src1.access_ptr("r")
-        src1_full_shape = list(src1.shape)
+        src1_full_shape = _fold_nd_shape_to_2d(src1.shape)
 
-    if len(src1_full_shape) == 1:
+    src0_buffer = src0.buffer if isinstance(src0, BufferRegion) else src0
+    src1_buffer = src1.buffer if isinstance(src1, BufferRegion) else src1
+    if DataType(src1_buffer.dtype) != DataType(src0_buffer.dtype):
+        mismatch = f"src1={src1_buffer.dtype}, src0={src0_buffer.dtype}"
+        raise ValueError(f"src1 and src0 dtypes must match: {mismatch}")
+    if isinstance(src1, BufferRegion):
+        _validate_buffer_region_flat_contiguity(src1)
+
+    dtype_bits = DataType(src0_buffer.dtype).bits
+    block_bits = 32 * 8
+    if block_bits % dtype_bits != 0:
+        raise ValueError(f"{op_name} does not support dtype {src0_buffer.dtype}.")
+    elems_per_block = block_bits // dtype_bits
+
+    if tmp is None:
+        if len(src1_full_shape) != 2:
+            requirement = f"packed src1 shape [R, {elems_per_block}] when tmp is omitted"
+            raise ValueError(f"{op_name} requires {requirement}; got {src1_full_shape}.")
+        s0, s1 = src1_full_shape[-2], src1_full_shape[-1]
+        src1_len = s0
+        if not _const_equal(s1, elems_per_block):
+            requirement = f"packed src1 shape [R, {elems_per_block}] when tmp is omitted"
+            raise ValueError(f"{op_name} requires {requirement}; got {src1_full_shape}.")
+    elif len(src1_full_shape) == 1:
         src1_len = src1_full_shape[0]
     elif len(src1_full_shape) == 2:
         s0, s1 = src1_full_shape[-2], src1_full_shape[-1]
         if _is_const_one(s0):
             src1_len = s1
-        elif _is_const_one(s1) or _const_equal(s0, dst_shape[0]):
+        elif _is_const_one(s1):
             src1_len = s0
         else:
-            raise ValueError(f"src1 must be 1D [R], [1, R], or [R, 1]; got {src1_full_shape}")
+            requirement = "scalar src1 shape [R], [1, R], or [R, 1] when tmp is provided"
+            raise ValueError(f"{op_name} requires {requirement}; got {src1_full_shape}.")
     else:
         raise ValueError(f"src1 must be 1D or 2D, got shape {src1_full_shape}")
 
@@ -2354,9 +2421,33 @@ def _row_expand_binop_experiment(
         raise ValueError(f"dst and src0 shapes must match: dst={dst_shape}, src0={src0_shape}")
 
     if not _const_equal(src1_len, dst_shape[0]):
-        raise ValueError(f"src1 scalar count must match dst rows: src1={src1_len}, dst[0]={dst_shape[0]}")
+        mismatch = f"src1={src1_len}, dst[0]={dst_shape[0]}"
+        raise ValueError(f"src1 scalar count must match dst rows: {mismatch}")
+
+    if tmp is not None:
+        tmp_buffer = tmp.buffer if isinstance(tmp, BufferRegion) else tmp
+        if DataType(tmp_buffer.dtype) != DataType(src0_buffer.dtype):
+            mismatch = f"tmp={tmp_buffer.dtype}, src0={src0_buffer.dtype}"
+            raise ValueError(f"tmp and src0 dtypes must match: {mismatch}")
+        if isinstance(tmp, BufferRegion):
+            _validate_buffer_region_flat_contiguity(tmp)
+            tmp_shape = [region.extent for region in tmp.region]
+        else:
+            tmp_shape = list(tmp.shape)
+        tmp_size = math.prod(tmp_shape)
+        expected_tmp_size = dst_shape[0] * elems_per_block
+        if not _const_equal(tmp_size, expected_tmp_size):
+            requirement = f"tmp must contain {expected_tmp_size} elements"
+            raise ValueError(f"{op_name} {requirement}; got {tmp_size}.")
 
     dtype = _dtype(src0)
+    row_bits = 256 * 8
+    expected_cols = row_bits // dtype_bits
+    if not _const_equal(dst_shape[1], expected_cols):
+        dtype_name = src0_buffer.dtype
+        requirement = f"a 256-byte last dimension ({expected_cols} {dtype_name} elements)"
+        raise ValueError(f"{op_name} requires {requirement}, got {dst_shape[1]}.")
+
     args = [
         f"{op_name}<{dtype}>",
         dst_ptr,
@@ -2387,6 +2478,10 @@ def row_expand_mul_experiment(
 
     AscendC: brcb(src1→tmp) + mul_mask(dst, src0, tmp).
     PTO: TROWEXPANDMUL_row_vec(dst, src0, src1).
+
+    Contiguous leading dimensions of dst/src0 are folded into rows; each row
+    must be 256 bytes. Without tmp, src1 is packed as one 32-byte block per
+    row. With tmp, src1 is scalar-linear and tmp supplies those packed blocks.
     """
     return _row_expand_binop_experiment(
         dst,
@@ -2409,6 +2504,10 @@ def row_expand_sub_experiment(
 
     AscendC: brcb(src1→tmp) + sub_mask(dst, src0, tmp).
     PTO: TROWEXPANDSUB_row_vec(dst, src0, src1).
+
+    Contiguous leading dimensions of dst/src0 are folded into rows; each row
+    must be 256 bytes. Without tmp, src1 is packed as one 32-byte block per
+    row. With tmp, src1 is scalar-linear and tmp supplies those packed blocks.
     """
     return _row_expand_binop_experiment(
         dst,
@@ -2431,6 +2530,10 @@ def row_expand_div_experiment(
 
     AscendC: brcb(src1→tmp) + div_mask(dst, src0, tmp).
     PTO: TROWEXPANDDIV_row_vec(dst, src0, src1).
+
+    Contiguous leading dimensions of dst/src0 are folded into rows; each row
+    must be 256 bytes. Without tmp, src1 is packed as one 32-byte block per
+    row. With tmp, src1 is scalar-linear and tmp supplies those packed blocks.
     """
     return _row_expand_binop_experiment(
         dst,