Add merge-axis TileLang element-wise templates

lib/TileOps/merge_axis.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+
+"""Shared helpers for full-axis merge-axis TileOp implementations."""
+
+import tilelang_dsl as pto
+
+
+def full_axis_constraint(**attrs):
+    """Match only full-axis tiles whose valid extents equal their allocation shape."""
+    matched_any_tile = False
+    for attr_name, shape in attrs.items():
+        if not attr_name.endswith("_shape") or attr_name.endswith("_valid_shape"):
+            continue
+        valid_shape = attrs.get(attr_name.replace("_shape", "_valid_shape"))
+        if valid_shape is None:
+            continue
+        matched_any_tile = True
+        if tuple(shape) != tuple(valid_shape):
+            return False
+    return matched_any_tile
+
+
+@pto.inline_proc
+def emit_unary_merge_axis(src: pto.Tile, dst: pto.Tile, compute):
+    """Lower a full-axis unary element-wise kernel through one linearized loop."""
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        in_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            result = compute(pto.vlds(in_tile, lane), mask)
+            pto.vsts(result, out_tile, lane, mask)
+
+
+@pto.inline_proc
+def emit_binary_merge_axis(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile, compute):
+    """Lower a full-axis binary element-wise kernel through one linearized loop."""
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src0, src1, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        lhs_tile,
+        rhs_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            result = compute(pto.vlds(lhs_tile, lane), pto.vlds(rhs_tile, lane), mask)
+            pto.vsts(result, out_tile, lane, mask)

lib/TileOps/tabs_template.py

@@ -9,6 +9,37 @@
 """TileLang DSL template for pto.tabs"""
 
 import tilelang_dsl as pto
+from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint
+
+
+@pto.vkernel(
+    target="a5",
+    op="pto.tabs",
+    constraints=[full_axis_constraint],
+    priority=100,
+    advanced=True,
+)
+def template_tabs_merge_axis(src: pto.Tile, dst: pto.Tile):
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        in_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            vec = pto.vlds(in_tile, lane)
+            result = pto.vabs(vec, mask)
+            pto.vsts(result, out_tile, lane, mask)
+    return
 
 
 @pto.vkernel(

lib/TileOps/tadd_template.py

@@ -13,6 +13,55 @@
 import tilelang_dsl as pto
 
 
+def _constraint_tadd_full_axis(
+    src0_shape=(),
+    src0_valid_shape=(),
+    src1_shape=(),
+    src1_valid_shape=(),
+    dst_shape=(),
+    dst_valid_shape=(),
+):
+    """Select the merged-axis fast path only for full-axis contiguous tiles."""
+    return (
+        src0_shape == src0_valid_shape
+        and src1_shape == src1_valid_shape
+        and dst_shape == dst_valid_shape
+    )
+
+
+@pto.vkernel(
+    target="a5",
+    op="pto.tadd",
+    constraints=[_constraint_tadd_full_axis],
+    priority=100,
+    advanced=True,
+)
+def template_tadd_merge_axis(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile):
+    """Merged-axis full-tile version aligned with TBinOps_1D_NoPostUpdate."""
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src0, src1, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        lhs_tile,
+        rhs_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            lhs = pto.vlds(lhs_tile, lane)
+            rhs = pto.vlds(rhs_tile, lane)
+            summed = pto.vadd(lhs, rhs, mask)
+            pto.vsts(summed, out_tile, lane, mask)
+    return
+
+
 @pto.vkernel(
     target="a5",
     op="pto.tadd"

lib/TileOps/tadds_template.py

@@ -11,6 +11,38 @@
 import sys
 from pathlib import Path
 import tilelang_dsl as pto
+from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint
+
+
+@pto.vkernel(
+    target="a5",
+    op="pto.tadds",
+    constraints=[full_axis_constraint],
+    priority=100,
+    advanced=True,
+)
+def template_tadds_merge_axis(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile):
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src, scalar, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        in_tile,
+        scalar_value,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            vec = pto.vlds(in_tile, lane)
+            result = pto.vadds(vec, scalar_value, mask)
+            pto.vsts(result, out_tile, lane, mask)
+    return
 
 
 @pto.vkernel(

lib/TileOps/tand_template.py

@@ -11,6 +11,39 @@
 import sys
 from pathlib import Path
 import tilelang_dsl as pto
+from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint
+
+
+@pto.vkernel(
+    target="a5",
+    op="pto.tand",
+    constraints=[full_axis_constraint],
+    priority=100,
+    advanced=True,
+)
+def template_tand_merge_axis(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile):
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+
+    with pto.strict_vecscope(dst, src0, src1, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        lhs_tile,
+        rhs_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            lhs = pto.vlds(lhs_tile, lane)
+            rhs = pto.vlds(rhs_tile, lane)
+            result = pto.vand(lhs, rhs, mask)
+            pto.vsts(result, out_tile, lane, mask)
+    return
 
 
 @pto.vkernel(

lib/TileOps/tands_template.py

@@ -19,6 +19,38 @@
 import sys
 from pathlib import Path
 import tilelang_dsl as pto
+from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint
+
+
+@pto.vkernel(
+    target="a5",
+    op="pto.tands",
+    constraints=[full_axis_constraint],
+    priority=100,
+    advanced=True,
+)
+def template_tands_merge_axis(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile):
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+    with pto.strict_vecscope(dst, src, scalar, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        in_tile,
+        scalar_value,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        scalar_vec = pto.vbr(scalar_value)
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            vec = pto.vlds(in_tile, lane)
+            result = pto.vand(vec, scalar_vec, mask)
+            pto.vsts(result, out_tile, lane, mask)
+    return
 
 
 @pto.vkernel(

lib/TileOps/tdiv_template.py

@@ -11,11 +11,50 @@
 import sys
 from pathlib import Path
 import tilelang_dsl as pto
+from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint
 
 # Import shared high-precision division algorithms
 from div_hp import _div_ieee754_f32_impl, _div_ieee754_f16_impl
 
 
+@pto.vkernel(
+    target="a5",
+    op="pto.tdiv",
+    constraints=[full_axis_constraint],
+    priority=100,
+    advanced=True,
+)
+def template_tdiv_merge_axis(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile):
+    dtype = dst.element_type
+    valid_rows, valid_cols = dst.valid_shape
+    total_elems = valid_rows * valid_cols
+    lanes = pto.get_lanes(dtype)
+    with pto.strict_vecscope(dst, src0, src1, total_elems, 0, total_elems, lanes) as (
+        out_tile,
+        lhs_tile,
+        rhs_tile,
+        area,
+        lb,
+        ub,
+        step,
+    ):
+        precision_mode = pto.get_op_attr("precision_mode", "DEFAULT")
+        remained = area
+        for lane in range(lb, ub, step):
+            mask, remained = pto.make_mask(out_tile.element_type, remained)
+            lhs = pto.vlds(lhs_tile, lane)
+            rhs = pto.vlds(rhs_tile, lane)
+            if pto.constexpr(precision_mode == "HIGH_PRECISION"):
+                if pto.constexpr(out_tile.element_type == pto.f32):
+                    result = _div_ieee754_f32_impl(lhs, rhs, mask)
+                else:
+                    result = _div_ieee754_f16_impl(lhs, rhs, mask)
+            else:
+                result = pto.vdiv(lhs, rhs, mask)
+            pto.vsts(result, out_tile, lane, mask)
+    return
+
+
 @pto.vkernel(
     target="a5",
     op="pto.tdiv"
@@ -47,4 +86,4 @@ def template_tdiv(src0: pto.Tile, src1: pto.Tile, dst: pto.Tile):
                 rhs = pto.vlds(src1[row, col:])
                 divided = pto.vdiv(lhs, rhs, mask)
                 pto.vsts(divided, dst[row, col:], mask)
-    return
+    return