From e9cfb6071a7f7a914d81e580f6a2bd2a8a1559d5 Mon Sep 17 00:00:00 2001 From: liuzidi Date: Thu, 28 May 2026 13:08:22 +0000 Subject: [PATCH] Add merge-axis TileLang element-wise templates --- lib/TileOps/merge_axis.py | 73 +++++++++++++++++ lib/TileOps/tabs_template.py | 31 ++++++++ lib/TileOps/tadd_template.py | 49 ++++++++++++ lib/TileOps/tadds_template.py | 32 ++++++++ lib/TileOps/tand_template.py | 33 ++++++++ lib/TileOps/tands_template.py | 32 ++++++++ lib/TileOps/tdiv_template.py | 41 +++++++++- lib/TileOps/tdivs_template.py | 79 ++++++++++++++++++- lib/TileOps/texp_template.py | 64 +++++++++++++++ lib/TileOps/tlog_template.py | 50 +++++++++++- lib/TileOps/tlrelu_template.py | 37 ++++++++- lib/TileOps/tmax_template.py | 33 ++++++++ lib/TileOps/tmaxs_template.py | 32 ++++++++ lib/TileOps/tmin_template.py | 33 ++++++++ lib/TileOps/tmins_template.py | 32 ++++++++ lib/TileOps/tmov_template.py | 31 ++++++++ lib/TileOps/tmul_template.py | 33 ++++++++ lib/TileOps/tmuls_template.py | 32 ++++++++ lib/TileOps/tneg_template.py | 31 ++++++++ lib/TileOps/tnot_template.py | 32 ++++++++ lib/TileOps/tor_template.py | 33 ++++++++ lib/TileOps/tors_template.py | 32 ++++++++ lib/TileOps/tprelu_template.py | 37 ++++++++- lib/TileOps/trecip_template.py | 59 +++++++++++++- lib/TileOps/trelu_template.py | 32 ++++++++ lib/TileOps/trsqrt_template.py | 51 +++++++++++- lib/TileOps/tshl_template.py | 33 ++++++++ lib/TileOps/tshls_template.py | 32 ++++++++ lib/TileOps/tshr_template.py | 33 ++++++++ lib/TileOps/tshrs_template.py | 32 ++++++++ lib/TileOps/tsqrt_template.py | 66 +++++++++++++++- lib/TileOps/tsub_template.py | 33 ++++++++ lib/TileOps/tsubs_template.py | 32 ++++++++ lib/TileOps/txor_template.py | 35 +++++++- lib/TileOps/txors_template.py | 32 ++++++++ .../src/st/testcase/run_ptoas_to_file.cmake | 12 +++ .../npu/a5/src/st/testcase/tabs/cases.py | 7 ++ .../npu/a5/src/st/testcase/tabs/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tabs/main.cpp | 3 + .../npu/a5/src/st/testcase/tabs/tabs.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tadd/cases.py | 7 ++ .../npu/a5/src/st/testcase/tadd/launch.cpp | 11 ++- .../npu/a5/src/st/testcase/tadd/main.cpp | 2 + .../npu/a5/src/st/testcase/tadd/tadd.pto | 63 ++++++++++++++- .../npu/a5/src/st/testcase/tadds/cases.py | 9 ++- .../npu/a5/src/st/testcase/tadds/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tadds/main.cpp | 3 + .../npu/a5/src/st/testcase/tadds/tadds.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tand/cases.py | 7 ++ .../npu/a5/src/st/testcase/tand/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tand/main.cpp | 5 +- .../npu/a5/src/st/testcase/tand/tand.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tands/cases.py | 9 ++- .../npu/a5/src/st/testcase/tands/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tands/main.cpp | 2 + .../npu/a5/src/st/testcase/tands/tands.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tdiv/cases.py | 8 ++ .../npu/a5/src/st/testcase/tdiv/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tdiv/main.cpp | 2 + .../npu/a5/src/st/testcase/tdiv/tdiv.pto | 58 ++++++++++++++ .../npu/a5/src/st/testcase/tdivs/cases.py | 8 ++ .../npu/a5/src/st/testcase/tdivs/launch.cpp | 8 +- .../npu/a5/src/st/testcase/tdivs/main.cpp | 2 + .../npu/a5/src/st/testcase/tdivs/tdivs.pto | 22 ++++++ .../npu/a5/src/st/testcase/texp/cases.py | 8 ++ .../npu/a5/src/st/testcase/texp/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/texp/main.cpp | 2 + .../npu/a5/src/st/testcase/texp/texp.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tlog/cases.py | 7 ++ .../npu/a5/src/st/testcase/tlog/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tlog/main.cpp | 2 + .../npu/a5/src/st/testcase/tlog/tlog.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tlrelu/cases.py | 11 ++- .../npu/a5/src/st/testcase/tlrelu/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tlrelu/main.cpp | 5 +- .../npu/a5/src/st/testcase/tlrelu/tlrelu.pto | 53 ++++++++++++- .../npu/a5/src/st/testcase/tmax/cases.py | 7 ++ .../npu/a5/src/st/testcase/tmax/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tmax/main.cpp | 2 + .../npu/a5/src/st/testcase/tmax/tmax.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tmaxs/cases.py | 3 +- .../npu/a5/src/st/testcase/tmaxs/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tmaxs/main.cpp | 2 + .../npu/a5/src/st/testcase/tmaxs/tmaxs.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tmin/cases.py | 7 ++ .../npu/a5/src/st/testcase/tmin/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tmin/main.cpp | 2 + .../npu/a5/src/st/testcase/tmin/tmin.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tmins/cases.py | 3 +- .../npu/a5/src/st/testcase/tmins/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tmins/main.cpp | 2 + .../npu/a5/src/st/testcase/tmins/tmins.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tmov/cases.py | 7 ++ .../npu/a5/src/st/testcase/tmov/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tmov/main.cpp | 3 + .../npu/a5/src/st/testcase/tmov/tmov.pto | 45 ++++++++++- .../npu/a5/src/st/testcase/tmul/cases.py | 7 ++ .../npu/a5/src/st/testcase/tmul/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tmul/main.cpp | 2 + .../npu/a5/src/st/testcase/tmul/tmul.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tmuls/cases.py | 9 ++- .../npu/a5/src/st/testcase/tmuls/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tmuls/main.cpp | 2 + .../npu/a5/src/st/testcase/tmuls/tmuls.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tneg/cases.py | 7 ++ .../npu/a5/src/st/testcase/tneg/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tneg/main.cpp | 3 + .../npu/a5/src/st/testcase/tneg/tneg.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tnot/cases.py | 9 ++- .../npu/a5/src/st/testcase/tnot/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tnot/main.cpp | 3 + .../npu/a5/src/st/testcase/tnot/tnot.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tor/cases.py | 7 ++ .../npu/a5/src/st/testcase/tor/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tor/main.cpp | 5 +- .../npu/a5/src/st/testcase/tor/tor.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tors/cases.py | 9 ++- .../npu/a5/src/st/testcase/tors/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tors/main.cpp | 2 + .../npu/a5/src/st/testcase/tors/tors.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tprelu/cases.py | 9 ++- .../npu/a5/src/st/testcase/tprelu/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tprelu/main.cpp | 5 +- .../npu/a5/src/st/testcase/tprelu/tprelu.pto | 63 ++++++++++++++- .../npu/a5/src/st/testcase/trecip/cases.py | 7 ++ .../npu/a5/src/st/testcase/trecip/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/trecip/main.cpp | 2 + .../npu/a5/src/st/testcase/trecip/trecip.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/trelu/cases.py | 7 ++ .../npu/a5/src/st/testcase/trelu/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/trelu/main.cpp | 5 +- .../npu/a5/src/st/testcase/trelu/trelu.pto | 47 ++++++++++- .../npu/a5/src/st/testcase/trsqrt/cases.py | 7 ++ .../npu/a5/src/st/testcase/trsqrt/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/trsqrt/main.cpp | 2 + .../npu/a5/src/st/testcase/trsqrt/trsqrt.pto | 45 +++++++++++ .../npu/a5/src/st/testcase/tshl/cases.py | 7 ++ .../npu/a5/src/st/testcase/tshl/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tshl/main.cpp | 5 +- .../npu/a5/src/st/testcase/tshl/tshl.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tshls/cases.py | 9 ++- .../npu/a5/src/st/testcase/tshls/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tshls/main.cpp | 2 + .../npu/a5/src/st/testcase/tshls/tshls.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tshr/cases.py | 7 ++ .../npu/a5/src/st/testcase/tshr/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tshr/main.cpp | 5 +- .../npu/a5/src/st/testcase/tshr/tshr.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tshrs/cases.py | 9 ++- .../npu/a5/src/st/testcase/tshrs/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tshrs/main.cpp | 2 + .../npu/a5/src/st/testcase/tshrs/tshrs.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/tsqrt/cases.py | 8 ++ .../npu/a5/src/st/testcase/tsqrt/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tsqrt/main.cpp | 2 + .../npu/a5/src/st/testcase/tsqrt/tsqrt.pto | 47 ++++++++++- .../npu/a5/src/st/testcase/tsub/cases.py | 7 ++ .../npu/a5/src/st/testcase/tsub/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/tsub/main.cpp | 2 + .../npu/a5/src/st/testcase/tsub/tsub.pto | 60 +++++++++++++- .../npu/a5/src/st/testcase/tsubs/cases.py | 3 +- .../npu/a5/src/st/testcase/tsubs/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/tsubs/main.cpp | 2 + .../npu/a5/src/st/testcase/tsubs/tsubs.pto | 43 ++++++++++ .../npu/a5/src/st/testcase/txor/cases.py | 7 ++ .../npu/a5/src/st/testcase/txor/launch.cpp | 9 ++- .../npu/a5/src/st/testcase/txor/main.cpp | 5 +- .../npu/a5/src/st/testcase/txor/txor.pto | 61 +++++++++++++- .../npu/a5/src/st/testcase/txors/cases.py | 9 ++- .../npu/a5/src/st/testcase/txors/launch.cpp | 7 ++ .../npu/a5/src/st/testcase/txors/main.cpp | 2 + .../npu/a5/src/st/testcase/txors/txors.pto | 46 +++++++++++ 172 files changed, 3643 insertions(+), 72 deletions(-) create mode 100644 lib/TileOps/merge_axis.py diff --git a/lib/TileOps/merge_axis.py b/lib/TileOps/merge_axis.py new file mode 100644 index 0000000000..32abc21b6a --- /dev/null +++ b/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) diff --git a/lib/TileOps/tabs_template.py b/lib/TileOps/tabs_template.py index 6c6802ae52..23e5a0e42e 100644 --- a/lib/TileOps/tabs_template.py +++ b/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( diff --git a/lib/TileOps/tadd_template.py b/lib/TileOps/tadd_template.py index 8e247fd732..7d48211d21 100644 --- a/lib/TileOps/tadd_template.py +++ b/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" diff --git a/lib/TileOps/tadds_template.py b/lib/TileOps/tadds_template.py index 7c3ddb06c7..ae478672dc 100644 --- a/lib/TileOps/tadds_template.py +++ b/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( diff --git a/lib/TileOps/tand_template.py b/lib/TileOps/tand_template.py index 6c14771976..601d89bb45 100644 --- a/lib/TileOps/tand_template.py +++ b/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( diff --git a/lib/TileOps/tands_template.py b/lib/TileOps/tands_template.py index 91258e7930..c6fd7b8f4f 100644 --- a/lib/TileOps/tands_template.py +++ b/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( diff --git a/lib/TileOps/tdiv_template.py b/lib/TileOps/tdiv_template.py index 2b841e13be..11dcd4d485 100644 --- a/lib/TileOps/tdiv_template.py +++ b/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 \ No newline at end of file + return diff --git a/lib/TileOps/tdivs_template.py b/lib/TileOps/tdivs_template.py index fa35a98560..84f5e08fff 100644 --- a/lib/TileOps/tdivs_template.py +++ b/lib/TileOps/tdivs_template.py @@ -19,11 +19,88 @@ 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.tdivs", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tdivs_merge_axis_tile_scalar(src: pto.Tile, scalar: pto.AnyType, dst: pto.Tile): + dtype = src.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) + 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) + vec = pto.vlds(in_tile, lane) + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + if pto.constexpr(out_tile.element_type == pto.f32): + result = _div_ieee754_f32_impl(vec, scalar_vec, mask) + else: + result = _div_ieee754_f16_impl(vec, scalar_vec, mask) + else: + result = pto.vdiv(vec, scalar_vec, mask) + pto.vsts(result, out_tile, lane, mask) + return + + +@pto.vkernel( + target="a5", + op="pto.tdivs", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tdivs_merge_axis_scalar_tile(scalar: pto.AnyType, src: pto.Tile, dst: pto.Tile): + dtype = src.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) + 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) + vec = pto.vlds(in_tile, lane) + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + if pto.constexpr(out_tile.element_type == pto.f32): + result = _div_ieee754_f32_impl(scalar_vec, vec, mask) + else: + result = _div_ieee754_f16_impl(scalar_vec, vec, mask) + else: + result = pto.vdiv(scalar_vec, vec, mask) + pto.vsts(result, out_tile, lane, mask) + return + + @pto.vkernel( target="a5", op="pto.tdivs", @@ -89,4 +166,4 @@ def template_tdivs_scalar_tile(scalar: pto.AnyType, src: pto.Tile, dst: pto.Tile scalar_vec = pto.vbr(scalar) result = pto.vdiv(scalar_vec, vec, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/texp_template.py b/lib/TileOps/texp_template.py index 609cb2fadf..f4c8c12432 100644 --- a/lib/TileOps/texp_template.py +++ b/lib/TileOps/texp_template.py @@ -9,6 +9,7 @@ """TileLang DSL template for pto.texp""" import tilelang_dsl as pto +from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint from exp_hp import _tl_exp_precision @pto.inline_proc @@ -39,6 +40,69 @@ def template_texp_impl(src: pto.Tile, dst: pto.Tile): pto.vsts(result, dst[row, col:], mask) return +@pto.inline_proc +def template_texp_merge_axis_hp_impl(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 = _tl_exp_precision(vec, mask, out_tile.element_type) + pto.vsts(result, out_tile, lane, mask) + return + + +@pto.inline_proc +def template_texp_merge_axis_impl(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) + result = pto.vexp(pto.vlds(in_tile, lane), mask) + pto.vsts(result, out_tile, lane, mask) + return + + +@pto.vkernel( + target="a5", + op="pto.texp", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_texp_merge_axis(src: pto.Tile, dst: pto.Tile): + hp_mode = pto.get_op_attr("precision_mode") + if pto.constexpr(hp_mode == "HIGH_PRECISION"): + template_texp_merge_axis_hp_impl(src, dst) + else: + template_texp_merge_axis_impl(src, dst) + return + + @pto.vkernel( target="a5", op="pto.texp" diff --git a/lib/TileOps/tlog_template.py b/lib/TileOps/tlog_template.py index 9b3dfabcd8..71730e9650 100644 --- a/lib/TileOps/tlog_template.py +++ b/lib/TileOps/tlog_template.py @@ -9,6 +9,7 @@ """TileLang DSL template for pto.tlog""" import tilelang_dsl as pto +from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint @pto.inline_proc @@ -49,6 +50,53 @@ def _tlog_default(src: pto.Tile, dst: pto.Tile, dtype, valid_rows, valid_cols): return None +@pto.vkernel( + target="a5", + op="pto.tlog", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tlog_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, + ): + precision_mode = pto.get_op_attr("precision_mode", "DEFAULT") + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + if pto.constexpr(out_tile.element_type == pto.f16): + subnormal_threshold = pto.f16("0x03FF") + mul_factor = pto.f16("0x6400") + compensation = pto.f16(-6.931471805599453094172) + elif pto.constexpr(out_tile.element_type == pto.f32): + subnormal_threshold = pto.f32("0x007FFFFF") + mul_factor = pto.f32("0x4B000000") + compensation = pto.f32(-15.9423851528787421) + 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) + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + cmp_mask = pto.vcmps(vec, subnormal_threshold, mask, pto.CmpMode.LT) + scaled = pto.vmuls(vec, mul_factor, mask) + selected_input = pto.vsel(scaled, vec, cmp_mask) + log_result = pto.vln(selected_input, mask) + compensated = pto.vadds(log_result, compensation, mask) + result = pto.vsel(compensated, log_result, cmp_mask) + else: + result = pto.vln(vec, mask) + pto.vsts(result, out_tile, lane, mask) + return + + @pto.vkernel( target="a5", op="pto.tlog", @@ -63,4 +111,4 @@ def template_tlog(src: pto.Tile, dst: pto.Tile): _tlog_high_precision(src, dst, dtype, valid_rows, valid_cols) else: _tlog_default(src, dst, dtype, valid_rows, valid_cols) - return \ No newline at end of file + return diff --git a/lib/TileOps/tlrelu_template.py b/lib/TileOps/tlrelu_template.py index 33087ca513..898d205166 100644 --- a/lib/TileOps/tlrelu_template.py +++ b/lib/TileOps/tlrelu_template.py @@ -11,6 +11,41 @@ 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.tlrelu", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tlrelu_merge_axis(src: pto.Tile, slope: pto.f32, 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, slope, total_elems, 0, total_elems, lanes) as ( + out_tile, + in_tile, + slope_value, + area, + lb, + ub, + step, + ): + if pto.constexpr(out_tile.element_type == pto.f16): + slope_scalar = pto.f16(slope_value) + else: + slope_scalar = slope_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.vlrelu(vec, slope_scalar, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( @@ -43,4 +78,4 @@ def template_tlrelu(src: pto.Tile, slope: pto.f32, dst: pto.Tile): src_vec = pto.vlds(src[row, col:]) result = pto.vlrelu(src_vec, slope_scalar, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/tmax_template.py b/lib/TileOps/tmax_template.py index 8831d3eef9..500a0809d8 100644 --- a/lib/TileOps/tmax_template.py +++ b/lib/TileOps/tmax_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.tmax", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmax_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.vmax(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tmaxs_template.py b/lib/TileOps/tmaxs_template.py index 5c9e409a3d..6ad3c68e29 100644 --- a/lib/TileOps/tmaxs_template.py +++ b/lib/TileOps/tmaxs_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.tmaxs", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmaxs_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.vmaxs(vec, scalar_value, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tmin_template.py b/lib/TileOps/tmin_template.py index 61664b14d1..6fc5764016 100644 --- a/lib/TileOps/tmin_template.py +++ b/lib/TileOps/tmin_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.tmin", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmin_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.vmin(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tmins_template.py b/lib/TileOps/tmins_template.py index bda0df5f92..36aa04a422 100644 --- a/lib/TileOps/tmins_template.py +++ b/lib/TileOps/tmins_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.tmins", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmins_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.vmins(vec, scalar_value, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tmov_template.py b/lib/TileOps/tmov_template.py index 4b1567f911..fede0f6e9b 100644 --- a/lib/TileOps/tmov_template.py +++ b/lib/TileOps/tmov_template.py @@ -17,6 +17,7 @@ """ import tilelang_dsl as pto +from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint def _tmov_ub2ub_nd2nd_constraint(src: pto.Tile, dst: pto.Tile) -> bool: @@ -60,6 +61,36 @@ def _tmov_ub2ub_nd2nd_constraint(src: pto.Tile, dst: pto.Tile) -> bool: return True +@pto.vkernel( + target="a5", + op="pto.tmov", + constraints=[_tmov_ub2ub_nd2nd_constraint, full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmov_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 = vec + pto.vsts(result, out_tile, lane, mask) + return None + + @pto.vkernel( target="a5", op="pto.tmov", diff --git a/lib/TileOps/tmul_template.py b/lib/TileOps/tmul_template.py index ae7adf44ef..0800a6e29b 100644 --- a/lib/TileOps/tmul_template.py +++ b/lib/TileOps/tmul_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.tmul", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmul_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.vmul(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tmuls_template.py b/lib/TileOps/tmuls_template.py index 8d02ea8264..2251ca9561 100644 --- a/lib/TileOps/tmuls_template.py +++ b/lib/TileOps/tmuls_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.tmuls", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tmuls_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.vmuls(vec, scalar_value, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tneg_template.py b/lib/TileOps/tneg_template.py index 8e10ce4cac..aed6f266cb 100644 --- a/lib/TileOps/tneg_template.py +++ b/lib/TileOps/tneg_template.py @@ -9,6 +9,37 @@ """TileLang DSL template for pto.tneg""" 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.tneg", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tneg_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.vneg(vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tnot_template.py b/lib/TileOps/tnot_template.py index f4728e8530..4b099dedaa 100644 --- a/lib/TileOps/tnot_template.py +++ b/lib/TileOps/tnot_template.py @@ -9,6 +9,38 @@ """TileLang DSL template for pto.tnot""" 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.tnot", + dtypes=[(pto.AnyInt, pto.AnyInt)], + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tnot_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.vnot(vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tor_template.py b/lib/TileOps/tor_template.py index e8be63d5e9..0a7a3ca2f7 100644 --- a/lib/TileOps/tor_template.py +++ b/lib/TileOps/tor_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.tor", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tor_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.vor(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tors_template.py b/lib/TileOps/tors_template.py index 542ee4044f..d979f4451e 100644 --- a/lib/TileOps/tors_template.py +++ b/lib/TileOps/tors_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.tors", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tors_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.vor(vec, scalar_vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tprelu_template.py b/lib/TileOps/tprelu_template.py index 45f83be102..04a53699dd 100644 --- a/lib/TileOps/tprelu_template.py +++ b/lib/TileOps/tprelu_template.py @@ -9,6 +9,41 @@ """TileLang DSL template for pto.tprelu (Parametric ReLU)""" 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.tprelu", + dtypes=[(pto.f16, pto.f16, pto.f16, pto.f16), (pto.f32, pto.f32, pto.f32, pto.f32), + (pto.f16, pto.f16, pto.i8, pto.f16), (pto.f32, pto.f32, pto.i8, pto.f32)], + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tprelu_merge_axis(src0: pto.Tile, src1: pto.Tile, tmp: 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.vprelu(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( @@ -41,4 +76,4 @@ def template_tprelu(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile vec1 = pto.vlds(src1[row, col:]) result = pto.vprelu(vec0, vec1, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/trecip_template.py b/lib/TileOps/trecip_template.py index 44947f3354..543516c62e 100644 --- a/lib/TileOps/trecip_template.py +++ b/lib/TileOps/trecip_template.py @@ -13,11 +13,68 @@ """ 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.trecip", + dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32)], + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_trecip_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) + if pto.constexpr(dtype == pto.f32): + with pto.strict_vecscope(dst, src, pto.f32(1.0), total_elems, 0, total_elems, lanes) as ( + out_tile, + in_tile, + one_value, + 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) + vec = pto.vlds(in_tile, lane) + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + result = _div_ieee754_f32_impl(pto.vbr(one_value), vec, mask) + else: + result = pto.vdiv(pto.vbr(one_value), vec, mask) + pto.vsts(result, out_tile, lane, mask) + else: + with pto.strict_vecscope(dst, src, pto.f16(1.0), total_elems, 0, total_elems, lanes) as ( + out_tile, + in_tile, + one_value, + 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) + vec = pto.vlds(in_tile, lane) + if pto.constexpr(precision_mode == "HIGH_PRECISION"): + result = _div_ieee754_f16_impl(pto.vbr(one_value), vec, mask) + else: + result = pto.vdiv(pto.vbr(one_value), vec, mask) + pto.vsts(result, out_tile, lane, mask) + return + + @pto.vkernel( target="a5", op="pto.trecip", @@ -58,4 +115,4 @@ def template_trecip(src: pto.Tile, dst: pto.Tile): one = pto.vbr(one_scalar) result = pto.vdiv(one, vinput, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/trelu_template.py b/lib/TileOps/trelu_template.py index 0cdd0e7eb3..9c1bd5f3c1 100644 --- a/lib/TileOps/trelu_template.py +++ b/lib/TileOps/trelu_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.trelu", + dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32), (pto.i32, pto.i32)], + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_trelu_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.vrelu(vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/trsqrt_template.py b/lib/TileOps/trsqrt_template.py index 87368adcaf..bf8ba3dd24 100644 --- a/lib/TileOps/trsqrt_template.py +++ b/lib/TileOps/trsqrt_template.py @@ -9,8 +9,57 @@ """TileLang DSL template for pto.trsqrt""" import tilelang_dsl as pto +from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint # TODO: Add implementation for HIGH_PRECISION type +@pto.vkernel( + target="a5", + op="pto.trsqrt", + dtypes=[(pto.f16, pto.f16), (pto.f32, pto.f32)], + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_trsqrt_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) + if pto.constexpr(dtype == pto.f32): + with pto.strict_vecscope(dst, src, pto.f32(1.0), total_elems, 0, total_elems, lanes) as ( + out_tile, + in_tile, + one_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.vdiv(pto.vbr(one_value), pto.vsqrt(vec, mask), mask) + pto.vsts(result, out_tile, lane, mask) + else: + with pto.strict_vecscope(dst, src, pto.f16(1.0), total_elems, 0, total_elems, lanes) as ( + out_tile, + in_tile, + one_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.vdiv(pto.vbr(one_value), pto.vsqrt(vec, mask), mask) + pto.vsts(result, out_tile, lane, mask) + return + + @pto.vkernel( target="a5", op="pto.trsqrt", @@ -33,4 +82,4 @@ def template_trsqrt(src: pto.Tile, dst: pto.Tile): sqrt_result = pto.vsqrt(vinput, mask) result = pto.vdiv(one, sqrt_result, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/tshl_template.py b/lib/TileOps/tshl_template.py index d236c8940d..7cdd18ba7a 100644 --- a/lib/TileOps/tshl_template.py +++ b/lib/TileOps/tshl_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.tshl", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tshl_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.vshl(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tshls_template.py b/lib/TileOps/tshls_template.py index def0b03537..727b1a0d9f 100644 --- a/lib/TileOps/tshls_template.py +++ b/lib/TileOps/tshls_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.tshls", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tshls_merge_axis(src: pto.Tile, scalar: pto.i16, 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.vshls(vec, scalar_value, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tshr_template.py b/lib/TileOps/tshr_template.py index f16ba9abe2..e7c5981962 100644 --- a/lib/TileOps/tshr_template.py +++ b/lib/TileOps/tshr_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.tshr", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tshr_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.vshr(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tshrs_template.py b/lib/TileOps/tshrs_template.py index 8366a638bd..64d8f46d03 100644 --- a/lib/TileOps/tshrs_template.py +++ b/lib/TileOps/tshrs_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.tshrs", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tshrs_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.vshrs(vec, scalar_value, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tsqrt_template.py b/lib/TileOps/tsqrt_template.py index 7381c50aa4..a75733ec69 100644 --- a/lib/TileOps/tsqrt_template.py +++ b/lib/TileOps/tsqrt_template.py @@ -9,6 +9,7 @@ """TileLang DSL template for pto.tsqrt""" import tilelang_dsl as pto +from merge_axis import emit_binary_merge_axis, emit_unary_merge_axis, full_axis_constraint from sqrt_hp import _tl_sqrt_precision @pto.inline_proc @@ -39,6 +40,69 @@ def template_tsqrt_impl(src: pto.Tile, dst: pto.Tile): pto.vsts(result, dst[row, col:], mask) return +@pto.inline_proc +def template_tsqrt_merge_axis_hp_impl(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 = _tl_sqrt_precision(vec, mask, out_tile.element_type) + pto.vsts(result, out_tile, lane, mask) + return + + +@pto.inline_proc +def template_tsqrt_merge_axis_impl(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) + result = pto.vsqrt(pto.vlds(in_tile, lane), mask) + pto.vsts(result, out_tile, lane, mask) + return + + +@pto.vkernel( + target="a5", + op="pto.tsqrt", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tsqrt_merge_axis(src: pto.Tile, dst: pto.Tile): + hp_mode = pto.get_op_attr("precision_mode") + if pto.constexpr(hp_mode == "HIGH_PRECISION"): + template_tsqrt_merge_axis_hp_impl(src, dst) + else: + template_tsqrt_merge_axis_impl(src, dst) + return + + @pto.vkernel( target="a5", op="pto.tsqrt" @@ -49,4 +113,4 @@ def template_tsqrt(src: pto.Tile, dst: pto.Tile): template_tsqrt_hp_impl(src, dst) else: template_tsqrt_impl(src, dst) - return \ No newline at end of file + return diff --git a/lib/TileOps/tsub_template.py b/lib/TileOps/tsub_template.py index 81d1b13ddd..005d49a539 100644 --- a/lib/TileOps/tsub_template.py +++ b/lib/TileOps/tsub_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.tsub", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tsub_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.vsub(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/tsubs_template.py b/lib/TileOps/tsubs_template.py index 84dc8bfbd0..941685b37e 100644 --- a/lib/TileOps/tsubs_template.py +++ b/lib/TileOps/tsubs_template.py @@ -17,6 +17,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.tsubs", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_tsubs_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.vsub(vec, scalar_vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/lib/TileOps/txor_template.py b/lib/TileOps/txor_template.py index d2ca4f1f73..345fdc3b8c 100644 --- a/lib/TileOps/txor_template.py +++ b/lib/TileOps/txor_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.txor", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_txor_merge_axis(src0: pto.Tile, src1: pto.Tile, tmp: 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.vxor(lhs, rhs, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( @@ -29,4 +62,4 @@ def template_txor(src0: pto.Tile, src1: pto.Tile, tmp: pto.Tile, dst: pto.Tile): rhs = pto.vlds(src1[row, col:]) result = pto.vxor(lhs, rhs, mask) pto.vsts(result, dst[row, col:], mask) - return \ No newline at end of file + return diff --git a/lib/TileOps/txors_template.py b/lib/TileOps/txors_template.py index c060209500..63a94a5697 100644 --- a/lib/TileOps/txors_template.py +++ b/lib/TileOps/txors_template.py @@ -18,6 +18,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.txors", + constraints=[full_axis_constraint], + priority=100, + advanced=True, +) +def template_txors_merge_axis(src: pto.Tile, scalar: pto.AnyType, tmp: 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, 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.vxor(vec, scalar_vec, mask) + pto.vsts(result, out_tile, lane, mask) + return @pto.vkernel( diff --git a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake index b8a3a0070a..9fbe978816 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake +++ b/test/tilelang_st/npu/a5/src/st/testcase/run_ptoas_to_file.cmake @@ -13,6 +13,14 @@ endif() get_filename_component(KERNEL_FATOBJ_DIR "${KERNEL_FATOBJ}" DIRECTORY) file(MAKE_DIRECTORY "${KERNEL_FATOBJ_DIR}") +get_filename_component( + PTOAS_REPO_ROOT + "${CMAKE_CURRENT_LIST_DIR}/../../../../../../.." + ABSOLUTE +) +set(PTOAS_TILELANG_PATH "${PTOAS_REPO_ROOT}/lib/TileOps") +set(PTOAS_TILELANG_PKG_PATH "${PTOAS_REPO_ROOT}/tilelang-dsl/python") + if(NOT DEFINED PTOAS_ENABLE_INSERT_SYNC) set(PTOAS_ENABLE_INSERT_SYNC ON) endif() @@ -27,6 +35,10 @@ if(DEFINED PTOAS_PTO_LEVEL AND NOT PTOAS_PTO_LEVEL STREQUAL "") endif() list(APPEND PTOAS_COMMAND --pto-backend=vpto) +list(APPEND PTOAS_COMMAND + "--tilelang-path=${PTOAS_TILELANG_PATH}" + "--tilelang-pkg-path=${PTOAS_TILELANG_PKG_PATH}" +) if(PTOAS_ENABLE_INSERT_SYNC) list(APPEND PTOAS_COMMAND --enable-insert-sync) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tabs/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tabs/cases.py index d63eb85f81..859a4a5575 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tabs/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tabs/cases.py @@ -52,4 +52,11 @@ "valid_shape": (32, 32), "eps": 1e-3, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tabs/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tabs/launch.cpp index dd39abd155..8e4c0023ee 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tabs/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tabs/launch.cpp @@ -38,4 +38,11 @@ extern "C" __global__ AICORE void TABS_f16_32x32(__gm__ uint16_t *a, __gm__ uint void LaunchTABS_f16_32x32(void *a, void *b, void *stream) { TABS_f16_32x32<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TABS_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTABS_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TABS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tabs/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tabs/main.cpp index 681510ddf3..07ea8c71b1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tabs/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tabs/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTABS_f32_16x64(void *a, void *b, void *stream); +void LaunchTABS_f32_merge_axis_16x72(void *a, void *b, void *stream); + void LaunchTABS_f32_32x32(void *a, void *b, void *stream); void LaunchTABS_f16_16x64(void *a, void *b, void *stream); void LaunchTABS_f16_32x32(void *a, void *b, void *stream); @@ -44,6 +46,7 @@ static const TestCase kCases[] = { {"f32_32x32", LaunchTABS_f32_32x32, 32, 32, 32, 32, sizeof(float)}, {"f16_16x64", LaunchTABS_f16_16x64, 16, 64, 16, 64, sizeof(uint16_t)}, {"f16_32x32", LaunchTABS_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", LaunchTABS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tabs/tabs.pto b/test/tilelang_st/npu/a5/src/st/testcase/tabs/tabs.pto index bf702ecdfe..b2956477f1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tabs/tabs.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tabs/tabs.pto @@ -177,4 +177,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TABS_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.tabs ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadd/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tadd/cases.py index 5958f05d2c..41c01dc372 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadd/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadd/cases.py @@ -23,6 +23,13 @@ import numpy as np CASES = [ + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, { "name": "f32_16x64", "dtype": np.float32, diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadd/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tadd/launch.cpp index f1074c8382..aa9e7e7a7a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadd/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadd/launch.cpp @@ -12,14 +12,21 @@ #define AICORE [aicore] #endif -// Case 0: f32 16x64 +// Case 0: f32 merge-axis full-axis 16x68 +extern "C" __global__ AICORE void TADD_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTADD_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream) { + TADD_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); +} + +// Case 1: f32 16x64 extern "C" __global__ AICORE void TADD_f32_16x64(__gm__ float *a, __gm__ float *b, __gm__ float *c); void LaunchTADD_f32_16x64(float *a, float *b, float *c, void *stream) { TADD_f32_16x64<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); } -// Case 1: f32 32x32 +// Case 2: f32 32x32 extern "C" __global__ AICORE void TADD_f32_32x32(__gm__ float *a, __gm__ float *b, __gm__ float *c); void LaunchTADD_f32_32x32(float *a, float *b, float *c, void *stream) { diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadd/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tadd/main.cpp index 1a010623f7..fbb499b879 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadd/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadd/main.cpp @@ -22,6 +22,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) +void LaunchTADD_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream); void LaunchTADD_f32_16x64(float *a, float *b, float *c, void *stream); void LaunchTADD_f32_32x32(float *a, float *b, float *c, void *stream); @@ -38,6 +39,7 @@ struct TestCase { }; static const TestCase kCases[] = { + {"f32_merge_axis_16x72", LaunchTADD_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, {"f32_16x64", LaunchTADD_f32_16x64, 16, 64, 16, 64, sizeof(float)}, {"f32_32x32", LaunchTADD_f32_32x32, 32, 32, 32, 32, sizeof(float)}, }; diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadd/tadd.pto b/test/tilelang_st/npu/a5/src/st/testcase/tadd/tadd.pto index 94d01af641..94f80c4844 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadd/tadd.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadd/tadd.pto @@ -12,8 +12,63 @@ // to produce a fatobj object. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - // Case 0: f32 16x64 (1024 elements) - func.func @TADD_f32_16x64(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.kernel} { + // Case 0: f32 merge-axis full-axis 16x72 (1152 elements) + func.func @TADD_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c72 = arith.constant 72 : index + %c1152 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c72], + strides = [%c1152, %c1152, %c1152, %c72, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c72], + strides = [%c1152, %c1152, %c1152, %c72, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c72], + strides = [%c1152, %c1152, %c1152, %c72, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c72] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c72] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c72] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tadd ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } + + // Case 1: f32 16x64 (1024 elements) + func.func @TADD_f32_16x64(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -67,8 +122,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.kernel} { + // Case 2: f32 32x32 (1024 elements) + func.func @TADD_f32_32x32(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c32 = arith.constant 32 : index diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadds/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tadds/cases.py index 5b24462bf4..3c7e0e1626 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadds/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadds/cases.py @@ -66,4 +66,11 @@ "valid_shape": (256, 16), "eps": 1e-6, }, -] + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadds/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tadds/launch.cpp index 49f0c98ecf..85008940ac 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadds/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadds/launch.cpp @@ -56,3 +56,10 @@ extern "C" __global__ AICORE void TADDS_f32_256x16(__gm__ float *src, __gm__ flo void LaunchTADDS_f32_256x16(float *src, float *dst, void *stream) { TADDS_f32_256x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TADDS_SCALAR_F32); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TADDS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); + +void LaunchTADDS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TADDS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TADDS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadds/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tadds/main.cpp index 4c6f409dc1..e7e93a3557 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadds/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadds/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTADDS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTADDS_f32_merge_axis_16x72(float *src, float *dst, void *stream); + void LaunchTADDS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTADDS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTADDS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -46,6 +48,7 @@ static const TestCase kCases[] = { {"i16_15x192", (void (*)(void*,void*,void*))LaunchTADDS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, {"f32_7x448", (void (*)(void*,void*,void*))LaunchTADDS_f32_7x448, 7, 448, 7, 448, sizeof(float)}, {"f32_256x16", (void (*)(void*,void*,void*))LaunchTADDS_f32_256x16, 256, 16, 256, 16, sizeof(float)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTADDS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tadds/tadds.pto b/test/tilelang_st/npu/a5/src/st/testcase/tadds/tadds.pto index 2057fc8df4..7b0c9e929e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tadds/tadds.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tadds/tadds.pto @@ -253,4 +253,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + pto.tadds ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tand/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tand/cases.py index 8c40489b19..d22ed07930 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tand/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tand/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 0, }, + { + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tand/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tand/launch.cpp index ed3149c6ee..b8f5fae851 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tand/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tand/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TAND_i32_32x32(__gm__ int32_t *a, __gm__ int32 void LaunchTAND_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream) { TAND_i32_32x32<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TAND_i32_merge_axis_16x72(__gm__ int32_t *a, __gm__ int32_t *b, __gm__ int32_t *c); + +void LaunchTAND_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream) { + TAND_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tand/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tand/main.cpp index 21b90e9b36..f44439cf6a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tand/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tand/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTAND_i32_16x64(int32_t *a, int32_t *b, int32_t *c, void *stream); +void LaunchTAND_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream); + void LaunchTAND_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream); using LaunchFn = void (*)(int32_t *, int32_t *, int32_t *, void *); @@ -40,6 +42,7 @@ struct TestCase { static const TestCase kCases[] = { {"i32_16x64", LaunchTAND_i32_16x64, 16, 64, 16, 64, sizeof(int32_t)}, {"i32_32x32", LaunchTAND_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, + {"i32_merge_axis_16x72", LaunchTAND_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -142,4 +145,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tand/tand.pto b/test/tilelang_st/npu/a5/src/st/testcase/tand/tand.pto index c1380f2b5a..40d017b7fd 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tand/tand.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tand/tand.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: i32 full-axis 16x72 (1152 elements) + func.func @TAND_i32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%b : !pto.tile_buf) + + pto.tand ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tands/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tands/cases.py index 18cc991784..5199832bf2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tands/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tands/cases.py @@ -39,4 +39,11 @@ "valid_shape": (15, 192), "eps": 0, }, -] +{ + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tands/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tands/launch.cpp index 8226ac79ea..16399fa70c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tands/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tands/launch.cpp @@ -43,3 +43,10 @@ extern "C" __global__ AICORE void TANDS_i16_15x192(__gm__ int16_t *src, __gm__ i void LaunchTANDS_i16_15x192(int16_t *src, int16_t *dst, void *stream) { TANDS_i16_15x192<<<1, nullptr, stream>>>((__gm__ int16_t *)src, (__gm__ int16_t *)dst, TANDS_SCALAR_I16); } + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TANDS_i32_merge_axis_16x72(__gm__ int32_t *src, __gm__ int32_t *dst, int32_t scalar); + +void LaunchTANDS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream) { + TANDS_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, TANDS_SCALAR_I32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tands/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tands/main.cpp index e0f93f2e73..ac0ca020ff 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tands/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tands/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTANDS_i32_32x64(int32_t *src, int32_t *dst, void *stream); +void LaunchTANDS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream); void LaunchTANDS_i16_63x64(int16_t *src, int16_t *dst, void *stream); void LaunchTANDS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTANDS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -42,6 +43,7 @@ static const TestCase kCases[] = { {"i16_63x64", (void (*)(void*,void*,void*))LaunchTANDS_i16_63x64, 63, 64, 63, 64, sizeof(int16_t)}, {"i32_31x128", (void (*)(void*,void*,void*))LaunchTANDS_i32_31x128, 31, 128, 31, 128, sizeof(int32_t)}, {"i16_15x192", (void (*)(void*,void*,void*))LaunchTANDS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, + {"i32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTANDS_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tands/tands.pto b/test/tilelang_st/npu/a5/src/st/testcase/tands/tands.pto index 24272369f9..41ee0adcde 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tands/tands.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tands/tands.pto @@ -173,4 +173,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: i32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + pto.tands ins(%src, %scalar : !pto.tile_buf, i32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/cases.py index 34e22f6338..10cb3220c3 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/cases.py @@ -203,4 +203,12 @@ "test_pattern": "precision_sensitive", "ulp_tolerance": 1, }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + "test_pattern": "normal", + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/launch.cpp index d4bbdb39ae..b38de9b53e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/launch.cpp @@ -131,3 +131,10 @@ extern "C" __global__ AICORE void TDIV_f16_2x32_hp(__gm__ void *a, __gm__ void * void LaunchTDIV_f16_2x32_hp(void *a, void *b, void *c, void *stream) { TDIV_f16_2x32_hp<<<1, nullptr, stream>>>(a, b, c); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TDIV_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTDIV_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream) { + TDIV_f32_merge_axis_16x72<<<1, nullptr, stream>>>(a, b, c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/main.cpp index c4f1a55d4e..60d4293927 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTDIV_f32_16x64(float *a, float *b, float *c, void *stream); +void LaunchTDIV_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream); void LaunchTDIV_f32_32x32(float *a, float *b, float *c, void *stream); void LaunchTDIV_f32_64x64(float *a, float *b, float *c, void *stream); void LaunchTDIV_f16_16x256(void *a, void *b, void *c, void *stream); @@ -71,6 +72,7 @@ static const TestCase kCases[] = { {"f16_16x64_hp_partial", (LaunchFn)LaunchTDIV_f16_16x64_hp_partial, 16, 64, 16, 63, 2}, {"f32_2x16_hp", (LaunchFn)LaunchTDIV_f32_2x16_hp, 2, 16, 2, 16, 4}, {"f16_2x32_hp", (LaunchFn)LaunchTDIV_f16_2x32_hp, 2, 32, 2, 32, 2}, + {"f32_merge_axis_16x72", (LaunchFn)LaunchTDIV_f32_merge_axis_16x72, 16, 72, 16, 72, 4}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/tdiv.pto b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/tdiv.pto index 3c4ad7ae59..f4cecf9ce5 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdiv/tdiv.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdiv/tdiv.pto @@ -968,4 +968,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TDIV_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tdiv ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/cases.py index c264d7805f..2890fa64e7 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/cases.py @@ -242,4 +242,12 @@ "direction": "scalar_src", "test_pattern": "overflow", }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + "direction": "src_scalar", + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/launch.cpp index 3b6cae07c4..8412a84340 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/launch.cpp @@ -148,4 +148,10 @@ void LaunchTDIVS_f32_16x64_hp_overflow_scalar_src(float *src, float *dst, void * extern "C" __global__ AICORE void TDIVS_f16_16x64_hp_overflow_scalar_src(__gm__ unsigned short *src, __gm__ unsigned short *dst, unsigned short scalar); void LaunchTDIVS_f16_16x64_hp_overflow_scalar_src(unsigned short *src, unsigned short *dst, void *stream) { TDIVS_f16_16x64_hp_overflow_scalar_src<<<1, nullptr, stream>>>((__gm__ unsigned short *)src, (__gm__ unsigned short *)dst, (unsigned short)0x7BEF); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TDIVS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); +void LaunchTDIVS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TDIVS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TDIVS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/main.cpp index 413cdc0f0d..c8ba5ce2f5 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTDIVS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTDIVS_f32_merge_axis_16x72(float *src, float *dst, void *stream); void LaunchTDIVS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTDIVS_f32_7x448(float *src, float *dst, void *stream); void LaunchTDIVS_f32_256x16(float *src, float *dst, void *stream); @@ -77,6 +78,7 @@ static const TestCase kCases[] = { {"f16_16x64_hp_subnormal_scalar_src", (void (*)(void*,void*,void*))LaunchTDIVS_f16_16x64_hp_subnormal_scalar_src, 16, 64, 16, 64, sizeof(uint16_t)}, {"f32_16x64_hp_overflow_scalar_src", (void (*)(void*,void*,void*))LaunchTDIVS_f32_16x64_hp_overflow_scalar_src, 16, 64, 16, 64, sizeof(float)}, {"f16_16x64_hp_overflow_scalar_src", (void (*)(void*,void*,void*))LaunchTDIVS_f16_16x64_hp_overflow_scalar_src, 16, 64, 16, 64, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTDIVS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/tdivs.pto b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/tdivs.pto index 0e70396b5e..364c846c1d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tdivs/tdivs.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tdivs/tdivs.pto @@ -396,4 +396,26 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + %src_view = pto.make_tensor_view %src_ptr, shape = [%c1, %c1, %c1, %c32, %c64], strides = [%c2048, %c2048, %c2048, %c64, %c1] : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, shape = [%c1, %c1, %c1, %c32, %c64], strides = [%c2048, %c2048, %c2048, %c64, %c1] : !pto.tensor_view<1x1x1x16x72xf32> + %src_part = pto.partition_view %src_view, offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c32, %c64] : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c32, %c64] : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) outs(%src : !pto.tile_buf) + pto.tdivs ins(%src, %scalar : !pto.tile_buf, f32) outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/texp/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/texp/cases.py index d303bf358e..80aaf4bdad 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/texp/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/texp/cases.py @@ -74,4 +74,12 @@ "eps": 1e-7, "high_precision": True, }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + "high_precision": False, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/texp/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/texp/launch.cpp index 59d9553132..6620774caf 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/texp/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/texp/launch.cpp @@ -52,4 +52,11 @@ extern "C" __global__ AICORE void TEXP_f16_64x64_hp2(__gm__ uint16_t *a, __gm__ void LaunchTEXP_f16_64x64_hp2(void *a, void *b, void *stream) { TEXP_f16_64x64_hp2<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TEXP_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTEXP_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TEXP_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/texp/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/texp/main.cpp index 7483e25516..b998dda648 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/texp/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/texp/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTEXP_f32_16x64(void *a, void *b, void *stream); +void LaunchTEXP_f32_merge_axis_16x72(void *a, void *b, void *stream); void LaunchTEXP_f32_32x32(void *a, void *b, void *stream); void LaunchTEXP_f16_16x64(void *a, void *b, void *stream); void LaunchTEXP_f16_32x32(void *a, void *b, void *stream); @@ -48,6 +49,7 @@ static const TestCase kCases[] = { {"f16_32x32", LaunchTEXP_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, {"f32_64x64_hp1", LaunchTEXP_f32_64x64_hp1, 64, 64, 64, 64, sizeof(float)}, {"f16_64x64_hp2", LaunchTEXP_f16_64x64_hp2, 64, 64, 64, 64, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", LaunchTEXP_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/texp/texp.pto b/test/tilelang_st/npu/a5/src/st/testcase/texp/texp.pto index 982a6dec5c..f7df60da7f 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/texp/texp.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/texp/texp.pto @@ -261,4 +261,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TEXP_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.texp ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlog/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tlog/cases.py index 1162c5d34b..3ffc84db4c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlog/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlog/cases.py @@ -84,4 +84,11 @@ "eps": 1e-3, "precision_mode": "HIGH_PRECISION", }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlog/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tlog/launch.cpp index 5fbe4521af..dd910551ab 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlog/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlog/launch.cpp @@ -66,4 +66,11 @@ extern "C" __global__ AICORE void TLOG_f16_32x32_hp(__gm__ uint16_t *a, __gm__ u void LaunchTLOG_f16_32x32_hp(void *a, void *b, void *stream) { TLOG_f16_32x32_hp<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TLOG_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTLOG_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TLOG_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlog/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tlog/main.cpp index b09c0a9486..a5e28be46a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlog/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlog/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTLOG_f32_16x64(void *a, void *b, void *stream); +void LaunchTLOG_f32_merge_axis_16x72(void *a, void *b, void *stream); void LaunchTLOG_f32_32x32(void *a, void *b, void *stream); void LaunchTLOG_f16_16x64(void *a, void *b, void *stream); void LaunchTLOG_f16_32x32(void *a, void *b, void *stream); @@ -52,6 +53,7 @@ static const TestCase kCases[] = { {"f32_32x32_hp", LaunchTLOG_f32_32x32_hp, 32, 32, 32, 32, sizeof(float)}, {"f16_16x64_hp", LaunchTLOG_f16_16x64_hp, 16, 64, 16, 64, sizeof(uint16_t)}, {"f16_32x32_hp", LaunchTLOG_f16_32x32_hp, 32, 32, 32, 32, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", LaunchTLOG_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlog/tlog.pto b/test/tilelang_st/npu/a5/src/st/testcase/tlog/tlog.pto index b2929fc291..e493f56769 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlog/tlog.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlog/tlog.pto @@ -347,4 +347,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TLOG_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.tlog ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/cases.py index a2b897e5d4..8a4618bd9b 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/cases.py @@ -62,4 +62,13 @@ "dst_valid_shape": (256, 16), "eps": 1e-3, }, -] \ No newline at end of file + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "dst_shape": (16, 72), + "dst_valid_shape": (16, 72), + "eps": 1e-3, + }, +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/launch.cpp index 356bd5115e..76fe22d8d4 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/launch.cpp @@ -38,4 +38,11 @@ extern "C" __global__ AICORE void TLRELU_f32_256x16_dst32(__gm__ float *src, __g void LaunchTLRELU_f32_256x16_dst32(float *src, float *dst, float slope, void *stream) { TLRELU_f32_256x16_dst32<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, slope); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TLRELU_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float slope); + +void LaunchTLRELU_f32_merge_axis_16x72(float *src, float *dst, float slope, void *stream) { + TLRELU_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, slope); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/main.cpp index 3b75edf69d..ca752f94b2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/main.cpp @@ -24,6 +24,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTLRELU_f32_32x64_dst128(float *src, float *dst, float slope, void *stream); +void LaunchTLRELU_f32_merge_axis_16x72(float *src, float *dst, float slope, void *stream); + void LaunchTLRELU_f16_63x64_dst128(uint16_t *src, uint16_t *dst, float slope, void *stream); void LaunchTLRELU_f32_7x448_dst512(float *src, float *dst, float slope, void *stream); void LaunchTLRELU_f32_256x16_dst32(float *src, float *dst, float slope, void *stream); @@ -48,6 +50,7 @@ static const TestCase kCases[] = { {"f16_63x64_dst128", (LaunchFn)LaunchTLRELU_f16_63x64_dst128, 63, 64, 63, 128, 63, 64, sizeof(uint16_t), true}, {"f32_7x448_dst512", (LaunchFn)LaunchTLRELU_f32_7x448_dst512, 7, 448, 7, 512, 7, 448, sizeof(float), false}, {"f32_256x16_dst32", (LaunchFn)LaunchTLRELU_f32_256x16_dst32, 256, 16, 256, 32, 256, 16, sizeof(float), false}, + {"f32_merge_axis_16x72", (LaunchFn)LaunchTLRELU_f32_merge_axis_16x72, 16, 72, 16, 72, 16, 72, sizeof(float), false}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -151,4 +154,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/tlrelu.pto b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/tlrelu.pto index bc44903975..75d0d05d0b 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/tlrelu.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tlrelu/tlrelu.pto @@ -211,4 +211,55 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TLRELU_f32_merge_axis_16x72(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %slope: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + // Src GM view: 1x1x1x16x72 + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + // Dst GM view: shape=valid_shape (16x72), strides based on dst allocation (16x72) + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + // Dst partition: sizes = valid_shape (16x72) + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + // Src UB tile: 16x72, valid 16x72 + %src_tile = pto.alloc_tile + : !pto.tile_buf + // Dst UB tile: 16x72, valid 16x72 + %dst_tile = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src_tile : !pto.tile_buf) + + pto.tlrelu ins(%src_tile, %slope : !pto.tile_buf, f32) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmax/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmax/cases.py index 69ba77ac4c..ebb78fc109 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmax/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmax/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 1e-6, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmax/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmax/launch.cpp index 3d47d685c4..cede58f2af 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmax/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmax/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TMAX_f32_32x32(__gm__ float *a, __gm__ float * void LaunchTMAX_f32_32x32(float *a, float *b, float *c, void *stream) { TMAX_f32_32x32<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMAX_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTMAX_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream) { + TMAX_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmax/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmax/main.cpp index 3dd9859a58..d93187edd4 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmax/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmax/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMAX_f32_16x64(float *a, float *b, float *c, void *stream); +void LaunchTMAX_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream); void LaunchTMAX_f32_32x32(float *a, float *b, float *c, void *stream); using LaunchFn = void (*)(float *, float *, float *, void *); @@ -40,6 +41,7 @@ struct TestCase { static const TestCase kCases[] = { {"f32_16x64", LaunchTMAX_f32_16x64, 16, 64, 16, 64, sizeof(float)}, {"f32_32x32", LaunchTMAX_f32_32x32, 32, 32, 32, 32, sizeof(float)}, + {"f32_merge_axis_16x72", LaunchTMAX_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmax/tmax.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmax/tmax.pto index 4a462fc2de..55f19003d5 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmax/tmax.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmax/tmax.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TMAX_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tmax ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/cases.py index d3bab221b1..9c84d4b3f8 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/cases.py @@ -19,4 +19,5 @@ {"name": "i16_15x192", "dtype": np.int16, "shape": (15, 192), "valid_shape": (15, 192), "eps": 0}, {"name": "f32_7x448", "dtype": np.float32, "shape": (7, 448), "valid_shape": (7, 448), "eps": 1e-6}, {"name": "f32_256x16", "dtype": np.float32, "shape": (256, 16), "valid_shape": (256, 16), "eps": 1e-6}, -] +{"name": "f32_merge_axis_16x72", "dtype": np.float32, "shape": (16, 72), "valid_shape": (16, 72), "eps": 1e-6}, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/launch.cpp index 793db13f1c..6a55e13895 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/launch.cpp @@ -56,3 +56,10 @@ extern "C" __global__ AICORE void TMAXS_f32_256x16(__gm__ float *src, __gm__ flo void LaunchTMAXS_f32_256x16(float *src, float *dst, void *stream) { TMAXS_f32_256x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMAXS_SCALAR_F32); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMAXS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); + +void LaunchTMAXS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TMAXS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMAXS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/main.cpp index 7104ff7ad1..44eacfd1db 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMAXS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTMAXS_f32_merge_axis_16x72(float *src, float *dst, void *stream); void LaunchTMAXS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTMAXS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTMAXS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -46,6 +47,7 @@ static const TestCase kCases[] = { {"i16_15x192", (void (*)(void*,void*,void*))LaunchTMAXS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, {"f32_7x448", (void (*)(void*,void*,void*))LaunchTMAXS_f32_7x448, 7, 448, 7, 448, sizeof(float)}, {"f32_256x16", (void (*)(void*,void*,void*))LaunchTMAXS_f32_256x16, 256, 16, 256, 16, sizeof(float)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTMAXS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/tmaxs.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/tmaxs.pto index 8adf6cb4c2..8e9faa3834 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/tmaxs.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmaxs/tmaxs.pto @@ -253,4 +253,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + pto.tmaxs ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmin/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmin/cases.py index 15bbb58ea1..aef469cb80 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmin/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmin/cases.py @@ -79,4 +79,11 @@ "valid_shape": (16, 200), "eps": 0, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmin/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmin/launch.cpp index 95247e512f..76ec472508 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmin/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmin/launch.cpp @@ -66,4 +66,11 @@ extern "C" __global__ AICORE void TMIN_i16_20x512_v16x200(__gm__ int16_t *a, __g void LaunchTMIN_i16_20x512_v16x200(void *a, void *b, void *c, void *stream) { TMIN_i16_20x512_v16x200<<<1, nullptr, stream>>>((__gm__ int16_t *)a, (__gm__ int16_t *)b, (__gm__ int16_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMIN_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTMIN_f32_merge_axis_16x72(void *a, void *b, void *c, void *stream) { + TMIN_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmin/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmin/main.cpp index 214a35a221..8f82a69863 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmin/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmin/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMIN_f32_64x64(void *a, void *b, void *c, void *stream); +void LaunchTMIN_f32_merge_axis_16x72(void *a, void *b, void *c, void *stream); void LaunchTMIN_i32_64x64(void *a, void *b, void *c, void *stream); void LaunchTMIN_i16_64x64(void *a, void *b, void *c, void *stream); void LaunchTMIN_f16_64x64(void *a, void *b, void *c, void *stream); @@ -52,6 +53,7 @@ static const TestCase kCases[] = { {"i32_64x64_v60x60", LaunchTMIN_i32_64x64_v60x60, 64, 64, 60, 60, sizeof(int32_t)}, {"f16_2x4096_v1x3600", LaunchTMIN_f16_2x4096_v1x3600, 2, 4096, 1, 3600, sizeof(uint16_t)}, {"i16_20x512_v16x200", LaunchTMIN_i16_20x512_v16x200, 20, 512, 16, 200, sizeof(int16_t)}, + {"f32_merge_axis_16x72", LaunchTMIN_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmin/tmin.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmin/tmin.pto index e7ba6d9b73..971d2e04d6 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmin/tmin.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmin/tmin.pto @@ -447,4 +447,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TMIN_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tmin ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmins/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmins/cases.py index 4526c4182f..2aacd763c8 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmins/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmins/cases.py @@ -19,4 +19,5 @@ {"name": "i16_15x192", "dtype": np.int16, "shape": (15, 192), "valid_shape": (15, 192), "eps": 0}, {"name": "f32_7x448", "dtype": np.float32, "shape": (7, 448), "valid_shape": (7, 448), "eps": 1e-6}, {"name": "f32_256x16", "dtype": np.float32, "shape": (256, 16), "valid_shape": (256, 16), "eps": 1e-6}, -] +{"name": "f32_merge_axis_16x72", "dtype": np.float32, "shape": (16, 72), "valid_shape": (16, 72), "eps": 1e-6}, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmins/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmins/launch.cpp index 65d44ffc4b..c59bc5870e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmins/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmins/launch.cpp @@ -56,3 +56,10 @@ extern "C" __global__ AICORE void TMINS_f32_256x16(__gm__ float *src, __gm__ flo void LaunchTMINS_f32_256x16(float *src, float *dst, void *stream) { TMINS_f32_256x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMINS_SCALAR_F32); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMINS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); + +void LaunchTMINS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TMINS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMINS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmins/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmins/main.cpp index 9fd09e48e7..dfdf6d32c8 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmins/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmins/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMINS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTMINS_f32_merge_axis_16x72(float *src, float *dst, void *stream); void LaunchTMINS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTMINS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTMINS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -46,6 +47,7 @@ static const TestCase kCases[] = { {"i16_15x192", (void (*)(void*,void*,void*))LaunchTMINS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, {"f32_7x448", (void (*)(void*,void*,void*))LaunchTMINS_f32_7x448, 7, 448, 7, 448, sizeof(float)}, {"f32_256x16", (void (*)(void*,void*,void*))LaunchTMINS_f32_256x16, 256, 16, 256, 16, sizeof(float)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTMINS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmins/tmins.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmins/tmins.pto index bb0f6b850e..36bc313138 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmins/tmins.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmins/tmins.pto @@ -253,4 +253,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + pto.tmins ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmov/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmov/cases.py index 3c6f4916c4..7a5670c80a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmov/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmov/cases.py @@ -95,4 +95,11 @@ "valid_shape": (128, 128), "eps": 0, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmov/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmov/launch.cpp index a88b8dbfc3..5974c17100 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmov/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmov/launch.cpp @@ -80,4 +80,11 @@ extern "C" __global__ AICORE void TMOV_u8_128x128(__gm__ uint8_t *src, __gm__ ui void LaunchTMOV_u8_128x128(uint8_t *src, uint8_t *dst, void *stream) { TMOV_u8_128x128<<<1, nullptr, stream>>>((__gm__ uint8_t *)src, (__gm__ uint8_t *)dst); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMOV_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst); + +void LaunchTMOV_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TMOV_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmov/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmov/main.cpp index bd1bfcac1b..f378ba61ef 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmov/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmov/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMOV_f32_64x64(float *src, float *dst, void *stream); +void LaunchTMOV_f32_merge_axis_16x72(float *src, float *dst, void *stream); + void LaunchTMOV_f32_32x32(float *src, float *dst, void *stream); void LaunchTMOV_f32_128x128(float *src, float *dst, void *stream); void LaunchTMOV_f32_128x32(float *src, float *dst, void *stream); @@ -54,6 +56,7 @@ static const TestCase kCases[] = { {"f16_128x128", (void(*)(void*,void*,void*))LaunchTMOV_f16_128x128, 128, 128, 128, 128, sizeof(uint16_t)}, {"u8_64x64", (void(*)(void*,void*,void*))LaunchTMOV_u8_64x64, 64, 64, 64, 64, sizeof(uint8_t)}, {"u8_128x128", (void(*)(void*,void*,void*))LaunchTMOV_u8_128x128, 128, 128, 128, 128, sizeof(uint8_t)}, + {"f32_merge_axis_16x72", (void(*)(void*,void*,void*))LaunchTMOV_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmov/tmov.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmov/tmov.pto index 8c91e55d3b..37655d792f 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmov/tmov.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmov/tmov.pto @@ -404,4 +404,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TMOV_f32_merge_axis_16x72(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmul/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmul/cases.py index 2d3a70ce83..01ffd0c317 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmul/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmul/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 1e-6, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmul/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmul/launch.cpp index 1debfe1406..1c6fc198d7 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmul/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmul/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TMUL_f32_32x32(__gm__ float *a, __gm__ float * void LaunchTMUL_f32_32x32(float *a, float *b, float *c, void *stream) { TMUL_f32_32x32<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMUL_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTMUL_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream) { + TMUL_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmul/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmul/main.cpp index 6e294af404..efdb578a8b 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmul/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmul/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMUL_f32_16x64(float *a, float *b, float *c, void *stream); +void LaunchTMUL_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream); void LaunchTMUL_f32_32x32(float *a, float *b, float *c, void *stream); using LaunchFn = void (*)(float *, float *, float *, void *); @@ -40,6 +41,7 @@ struct TestCase { static const TestCase kCases[] = { {"f32_16x64", LaunchTMUL_f32_16x64, 16, 64, 16, 64, sizeof(float)}, {"f32_32x32", LaunchTMUL_f32_32x32, 32, 32, 32, 32, sizeof(float)}, + {"f32_merge_axis_16x72", LaunchTMUL_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmul/tmul.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmul/tmul.pto index 7ee883cc94..ad5723ab4d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmul/tmul.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmul/tmul.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TMUL_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tmul ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/cases.py index d12724e5f1..31a422dbbf 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/cases.py @@ -66,4 +66,11 @@ "valid_shape": (256, 16), "eps": 1e-6, }, -] +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/launch.cpp index fdd67d5967..8b4a3e96db 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/launch.cpp @@ -56,3 +56,10 @@ extern "C" __global__ AICORE void TMULS_f32_256x16(__gm__ float *src, __gm__ flo void LaunchTMULS_f32_256x16(float *src, float *dst, void *stream) { TMULS_f32_256x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMULS_SCALAR_F32); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TMULS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); + +void LaunchTMULS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TMULS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TMULS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/main.cpp index a5372cccc5..e86cec46af 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTMULS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTMULS_f32_merge_axis_16x72(float *src, float *dst, void *stream); void LaunchTMULS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTMULS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTMULS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -46,6 +47,7 @@ static const TestCase kCases[] = { {"i16_15x192", (void (*)(void*,void*,void*))LaunchTMULS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, {"f32_7x448", (void (*)(void*,void*,void*))LaunchTMULS_f32_7x448, 7, 448, 7, 448, sizeof(float)}, {"f32_256x16", (void (*)(void*,void*,void*))LaunchTMULS_f32_256x16, 256, 16, 256, 16, sizeof(float)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTMULS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/tmuls.pto b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/tmuls.pto index a6fcc5cd9b..eb0258a44a 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tmuls/tmuls.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tmuls/tmuls.pto @@ -253,4 +253,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + pto.tmuls ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tneg/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tneg/cases.py index f5251d28a1..49d5a85ea0 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tneg/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tneg/cases.py @@ -66,4 +66,11 @@ "valid_shape": (64, 16), "eps": 0, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tneg/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tneg/launch.cpp index ef6121f48e..15d0c35e6c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tneg/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tneg/launch.cpp @@ -52,4 +52,11 @@ extern "C" __global__ AICORE void TNEG_i16_64x16(__gm__ int16_t *a, __gm__ int16 void LaunchTNEG_i16_64x16(void *a, void *b, void *stream) { TNEG_i16_64x16<<<1, nullptr, stream>>>((__gm__ int16_t *)a, (__gm__ int16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TNEG_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTNEG_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TNEG_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tneg/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tneg/main.cpp index 6a86073fc8..52ca8c5346 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tneg/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tneg/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTNEG_f32_16x64(void *a, void *b, void *stream); +void LaunchTNEG_f32_merge_axis_16x72(void *a, void *b, void *stream); + void LaunchTNEG_f32_32x32(void *a, void *b, void *stream); void LaunchTNEG_f16_16x64(void *a, void *b, void *stream); void LaunchTNEG_f16_32x32(void *a, void *b, void *stream); @@ -48,6 +50,7 @@ static const TestCase kCases[] = { {"f16_32x32", LaunchTNEG_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, {"i32_32x32", LaunchTNEG_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, {"i16_64x16", LaunchTNEG_i16_64x16, 64, 16, 64, 16, sizeof(int16_t)}, + {"f32_merge_axis_16x72", LaunchTNEG_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tneg/tneg.pto b/test/tilelang_st/npu/a5/src/st/testcase/tneg/tneg.pto index f67ec4f195..42bfe13766 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tneg/tneg.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tneg/tneg.pto @@ -260,4 +260,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TNEG_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.tneg ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tnot/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tnot/cases.py index b6612d63fe..204472d703 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tnot/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tnot/cases.py @@ -66,4 +66,11 @@ "valid_shape": (60, 60), "eps": 0 }, -] \ No newline at end of file + { + "name": "int8_merge_axis_32x288", + "dtype": np.int8, + "shape": (32, 288), + "valid_shape": (32, 288), + "eps": 0, + }, +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tnot/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tnot/launch.cpp index 858f6d1816..33acac8cab 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tnot/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tnot/launch.cpp @@ -52,4 +52,11 @@ extern "C" __global__ AICORE void TNOT_uint32_60x60(__gm__ uint32_t *a, __gm__ u void LaunchTNOT_uint32_60x60(void *a, void *b, void *stream) { TNOT_uint32_60x60<<<1, nullptr, stream>>>((__gm__ uint32_t *)a, (__gm__ uint32_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: int8 merge axis 32x264 +extern "C" __global__ AICORE void TNOT_int8_merge_axis_32x288(__gm__ int8_t *a, __gm__ int8_t *b); + +void LaunchTNOT_int8_merge_axis_32x288(void *a, void *b, void *stream) { + TNOT_int8_merge_axis_32x288<<<1, nullptr, stream>>>((__gm__ int8_t *)a, (__gm__ int8_t *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tnot/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tnot/main.cpp index 55a823be7b..3488855035 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tnot/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tnot/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTNOT_int8_64x64(void *a, void *b, void *stream); +void LaunchTNOT_int8_merge_axis_32x288(void *a, void *b, void *stream); + void LaunchTNOT_uint8_60x60(void *a, void *b, void *stream); void LaunchTNOT_int16_64x64(void *a, void *b, void *stream); void LaunchTNOT_uint16_60x60(void *a, void *b, void *stream); @@ -48,6 +50,7 @@ static const TestCase kCases[] = { {"uint16_60x60", LaunchTNOT_uint16_60x60, 64, 64, 60, 60, sizeof(uint16_t)}, {"int32_64x64", LaunchTNOT_int32_64x64, 64, 64, 64, 64, sizeof(int32_t)}, {"uint32_60x60", LaunchTNOT_uint32_60x60, 64, 64, 60, 60, sizeof(uint32_t)}, + {"int8_merge_axis_32x288", LaunchTNOT_int8_merge_axis_32x288, 32, 288, 32, 288, sizeof(int8_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tnot/tnot.pto b/test/tilelang_st/npu/a5/src/st/testcase/tnot/tnot.pto index bbd0288353..9a71a7c0de 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tnot/tnot.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tnot/tnot.pto @@ -260,4 +260,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: int8 full-axis 32x288 (9216 elements) + func.func @TNOT_int8_merge_axis_32x288(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %c288 = arith.constant 288 : index + %c9216 = arith.constant 9216 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c32, %c288], + strides = [%c9216, %c9216, %c9216, %c288, %c1] + : !pto.tensor_view<1x1x1x32x288xi8> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c32, %c288], + strides = [%c9216, %c9216, %c9216, %c288, %c1] + : !pto.tensor_view<1x1x1x32x288xi8> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c288] + : !pto.tensor_view<1x1x1x32x288xi8> -> !pto.partition_tensor_view<1x1x1x32x288xi8> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c288] + : !pto.tensor_view<1x1x1x32x288xi8> -> !pto.partition_tensor_view<1x1x1x32x288xi8> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x32x288xi8>) + outs(%a : !pto.tile_buf) + + pto.tnot ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x32x288xi8>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tor/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tor/cases.py index 736a5ff8ff..c995cc1618 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tor/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tor/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 0, }, + { + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tor/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tor/launch.cpp index 1cb9c1454f..415c258990 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tor/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tor/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TOR_i32_32x32(__gm__ int32_t *a, __gm__ int32_ void LaunchTOR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream) { TOR_i32_32x32<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TOR_i32_merge_axis_16x72(__gm__ int32_t *a, __gm__ int32_t *b, __gm__ int32_t *c); + +void LaunchTOR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream) { + TOR_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tor/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tor/main.cpp index 21d82eeea8..d1464cedb1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tor/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tor/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTOR_i32_16x64(int32_t *a, int32_t *b, int32_t *c, void *stream); +void LaunchTOR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream); + void LaunchTOR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream); using LaunchFn = void (*)(int32_t *, int32_t *, int32_t *, void *); @@ -40,6 +42,7 @@ struct TestCase { static const TestCase kCases[] = { {"i32_16x64", LaunchTOR_i32_16x64, 16, 64, 16, 64, sizeof(int32_t)}, {"i32_32x32", LaunchTOR_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, + {"i32_merge_axis_16x72", LaunchTOR_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -142,4 +145,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tor/tor.pto b/test/tilelang_st/npu/a5/src/st/testcase/tor/tor.pto index 0d6f7d9b31..1688caf3c3 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tor/tor.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tor/tor.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: i32 full-axis 16x72 (1152 elements) + func.func @TOR_i32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%b : !pto.tile_buf) + + pto.tor ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tors/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tors/cases.py index 18cc991784..5199832bf2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tors/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tors/cases.py @@ -39,4 +39,11 @@ "valid_shape": (15, 192), "eps": 0, }, -] +{ + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tors/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tors/launch.cpp index 4495ff38c0..ff0c26910c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tors/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tors/launch.cpp @@ -43,3 +43,10 @@ extern "C" __global__ AICORE void TORS_i16_15x192(__gm__ int16_t *src, __gm__ in void LaunchTORS_i16_15x192(int16_t *src, int16_t *dst, void *stream) { TORS_i16_15x192<<<1, nullptr, stream>>>((__gm__ int16_t *)src, (__gm__ int16_t *)dst, TORS_SCALAR_I16); } + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TORS_i32_merge_axis_16x72(__gm__ int32_t *src, __gm__ int32_t *dst, int32_t scalar); + +void LaunchTORS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream) { + TORS_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, TORS_SCALAR_I32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tors/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tors/main.cpp index b67da6f067..0adeb96ab6 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tors/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tors/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTORS_i32_32x64(int32_t *src, int32_t *dst, void *stream); +void LaunchTORS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream); void LaunchTORS_i16_63x64(int16_t *src, int16_t *dst, void *stream); void LaunchTORS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTORS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -42,6 +43,7 @@ static const TestCase kCases[] = { {"i16_63x64", (void (*)(void*,void*,void*))LaunchTORS_i16_63x64, 63, 64, 63, 64, sizeof(int16_t)}, {"i32_31x128", (void (*)(void*,void*,void*))LaunchTORS_i32_31x128, 31, 128, 31, 128, sizeof(int32_t)}, {"i16_15x192", (void (*)(void*,void*,void*))LaunchTORS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, + {"i32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTORS_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tors/tors.pto b/test/tilelang_st/npu/a5/src/st/testcase/tors/tors.pto index 36124ff380..2cfa40d960 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tors/tors.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tors/tors.pto @@ -173,4 +173,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: i32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + pto.tors ins(%src, %scalar : !pto.tile_buf, i32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/cases.py index bafcc28ef9..b3152f0027 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/cases.py @@ -80,4 +80,11 @@ "valid_shape": (2048, 8), "eps": 1e-6, }, -] \ No newline at end of file + { + "name": "f16_merge_axis_16x144", + "dtype": np.float16, + "shape": (16, 144), + "valid_shape": (16, 144), + "eps": 1e-3, + }, +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/launch.cpp index 4e15c68abb..d7e71ba372 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/launch.cpp @@ -66,4 +66,11 @@ extern "C" __global__ AICORE void TPRELU_f32_2048x8(__gm__ float *src0, __gm__ f void LaunchTPRELU_f32_2048x8(float *src0, float *src1, float *dst, void *stream) { TPRELU_f32_2048x8<<<1, nullptr, stream>>>((__gm__ float *)src0, (__gm__ float *)src1, (__gm__ float *)dst); -} \ No newline at end of file +} + +// Merge-axis case: f16 merge axis 16x136 +extern "C" __global__ AICORE void TPRELU_f16_merge_axis_16x144(__gm__ uint16_t *src0, __gm__ uint16_t *src1, __gm__ uint16_t *dst); + +void LaunchTPRELU_f16_merge_axis_16x144(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream) { + TPRELU_f16_merge_axis_16x144<<<1, nullptr, stream>>>((__gm__ uint16_t *)src0, (__gm__ uint16_t *)src1, (__gm__ uint16_t *)dst); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/main.cpp index a55be7945b..c4996bdaf1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTPRELU_f16_64x64(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream); +void LaunchTPRELU_f16_merge_axis_16x144(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream); + void LaunchTPRELU_f16_63x63(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream); void LaunchTPRELU_f16_1x16384(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream); void LaunchTPRELU_f16_2048x16(uint16_t *src0, uint16_t *src1, uint16_t *dst, void *stream); @@ -52,6 +54,7 @@ static const TestCase kCases[] = { {"f32_63x63", F32, (void*)LaunchTPRELU_f32_63x63, 64, 64, 63, 63}, {"f32_1x16384", F32, (void*)LaunchTPRELU_f32_1x16384, 1, 16384, 1, 16384}, {"f32_2048x8", F32, (void*)LaunchTPRELU_f32_2048x8, 2048, 8, 2048, 8}, + {"f16_merge_axis_16x144", F16, (void*)LaunchTPRELU_f16_merge_axis_16x144, 16, 144, 16, 144}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -195,4 +198,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/tprelu.pto b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/tprelu.pto index d72ed41a49..8528beabdc 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tprelu/tprelu.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tprelu/tprelu.pto @@ -491,4 +491,65 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f16 full-axis 16x136 (2176 elements) + func.func @TPRELU_f16_merge_axis_16x144(%src0_ptr: !pto.ptr, %src1_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c144 = arith.constant 144 : index + %c2304 = arith.constant 2304 : index + + %src0_view = pto.make_tensor_view %src0_ptr, + shape = [%c1, %c1, %c1, %c16, %c144], + strides = [%c2304, %c2304, %c2304, %c144, %c1] + : !pto.tensor_view<1x1x1x16x144xf16> + %src1_view = pto.make_tensor_view %src1_ptr, + shape = [%c1, %c1, %c1, %c16, %c144], + strides = [%c2304, %c2304, %c2304, %c144, %c1] + : !pto.tensor_view<1x1x1x16x144xf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c16, %c144], + strides = [%c2304, %c2304, %c2304, %c144, %c1] + : !pto.tensor_view<1x1x1x16x144xf16> + + %src0_part = pto.partition_view %src0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c144] + : !pto.tensor_view<1x1x1x16x144xf16> -> !pto.partition_tensor_view<1x1x1x16x144xf16> + %src1_part = pto.partition_view %src1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c144] + : !pto.tensor_view<1x1x1x16x144xf16> -> !pto.partition_tensor_view<1x1x1x16x144xf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c144] + : !pto.tensor_view<1x1x1x16x144xf16> -> !pto.partition_tensor_view<1x1x1x16x144xf16> + + %src0_tile = pto.alloc_tile + : !pto.tile_buf + %src1_tile = pto.alloc_tile + : !pto.tile_buf + %tmp_tile = pto.alloc_tile + : !pto.tile_buf + %dst_tile = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src0_part : !pto.partition_tensor_view<1x1x1x16x144xf16>) + outs(%src0_tile : !pto.tile_buf) + pto.tload ins(%src1_part : !pto.partition_tensor_view<1x1x1x16x144xf16>) + outs(%src1_tile : !pto.tile_buf) + + pto.tprelu ins(%src0_tile, %src1_tile, %tmp_tile : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x144xf16>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trecip/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/trecip/cases.py index b1c2012e28..4512948883 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trecip/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/trecip/cases.py @@ -66,4 +66,11 @@ "valid_shape": (58, 70), "eps": 1e-6, }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trecip/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trecip/launch.cpp index 3cf95d119c..3fbd72d800 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trecip/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trecip/launch.cpp @@ -52,4 +52,11 @@ extern "C" __global__ AICORE void TRECIP_f32_58x70(__gm__ float *a, __gm__ float void LaunchTRECIP_f32_58x70(void *a, void *b, void *stream) { TRECIP_f32_58x70<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TRECIP_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTRECIP_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TRECIP_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trecip/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trecip/main.cpp index 4a8400e6a3..a79f34716d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trecip/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trecip/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTRECIP_f32_16x64(void *a, void *b, void *stream); +void LaunchTRECIP_f32_merge_axis_16x72(void *a, void *b, void *stream); void LaunchTRECIP_f32_32x32(void *a, void *b, void *stream); void LaunchTRECIP_f16_16x64(void *a, void *b, void *stream); void LaunchTRECIP_f16_32x32(void *a, void *b, void *stream); @@ -48,6 +49,7 @@ static const TestCase kCases[] = { {"f16_32x32", LaunchTRECIP_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, {"f32_64x64_pad", LaunchTRECIP_f32_64x64_pad, 66, 72, 64, 64, sizeof(float)}, {"f32_58x70", LaunchTRECIP_f32_58x70, 66, 72, 58, 70, sizeof(float)}, + {"f32_merge_axis_16x72", LaunchTRECIP_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trecip/trecip.pto b/test/tilelang_st/npu/a5/src/st/testcase/trecip/trecip.pto index a2854aefa0..380b8348d1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trecip/trecip.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/trecip/trecip.pto @@ -265,4 +265,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TRECIP_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.trecip ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trelu/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/trelu/cases.py index 85823525c7..94fb45581c 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trelu/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/trelu/cases.py @@ -49,4 +49,11 @@ "valid_shape": (60, 60), "eps": 1e-6, }, + { + "name": "int32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trelu/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trelu/launch.cpp index 94e256ff4a..8de74351be 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trelu/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trelu/launch.cpp @@ -31,4 +31,11 @@ extern "C" __global__ AICORE void TRELU_f32_64x64_v60x60(__gm__ float *input, __ void LaunchTRELU_f32_64x64_v60x60(float *input, float *output, void *stream) { TRELU_f32_64x64_v60x60<<<1, nullptr, stream>>>((__gm__ float *)input, (__gm__ float *)output); -} \ No newline at end of file +} + +// Merge-axis case: int32 merge axis 16x72 +extern "C" __global__ AICORE void TRELU_int32_merge_axis_16x72(__gm__ int32_t *input, __gm__ int32_t *output); + +void LaunchTRELU_int32_merge_axis_16x72(int32_t *input, int32_t *output, void *stream) { + TRELU_int32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)input, (__gm__ int32_t *)output); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trelu/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trelu/main.cpp index b19a7fd95d..e5549d6a64 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trelu/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trelu/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTRELU_int32_64x64(int32_t *input, int32_t *output, void *stream); +void LaunchTRELU_int32_merge_axis_16x72(int32_t *input, int32_t *output, void *stream); + void LaunchTRELU_f16_64x64_v60x60(uint16_t *input, uint16_t *output, void *stream); void LaunchTRELU_f32_64x64_v60x60(float *input, float *output, void *stream); @@ -38,6 +40,7 @@ static const TestCase kCases[] = { {"int32_64x64", (void (*)(void*, void*, void*))LaunchTRELU_int32_64x64, 64, 64, sizeof(int32_t)}, {"f16_64x64_valid_60x60", (void (*)(void*, void*, void*))LaunchTRELU_f16_64x64_v60x60, 60, 60, sizeof(uint16_t)}, {"f32_64x64_valid_60x60", (void (*)(void*, void*, void*))LaunchTRELU_f32_64x64_v60x60, 60, 60, sizeof(float)}, + {"int32_merge_axis_16x72", (void (*)(void*, void*, void*))LaunchTRELU_int32_merge_axis_16x72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -125,4 +128,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trelu/trelu.pto b/test/tilelang_st/npu/a5/src/st/testcase/trelu/trelu.pto index 99c550b1dc..387bcdafb8 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trelu/trelu.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/trelu/trelu.pto @@ -136,4 +136,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: int32 full-axis 16x72 (1152 elements) + func.func @TRELU_int32_merge_axis_16x72(%input_ptr: !pto.ptr, %output_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %input_view = pto.make_tensor_view %input_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %output_view = pto.make_tensor_view %output_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %input_part = pto.partition_view %input_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %output_part = pto.partition_view %output_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%input_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + + pto.trelu ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + + pto.tstore ins(%dst : !pto.tile_buf) + outs(%output_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/cases.py index cb8b6a48d0..aab23c0193 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/cases.py @@ -52,4 +52,11 @@ "valid_shape": (32, 32), "eps": 1e-3, }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/launch.cpp index 65a35f3bd4..fb9e284a2d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/launch.cpp @@ -38,4 +38,11 @@ extern "C" __global__ AICORE void TRSQRT_f16_32x32(__gm__ uint16_t *a, __gm__ ui void LaunchTRSQRT_f16_32x32(void *a, void *b, void *stream) { TRSQRT_f16_32x32<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TRSQRT_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTRSQRT_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TRSQRT_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/main.cpp index 20c9550700..e1a915a8aa 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTRSQRT_f32_16x64(void *a, void *b, void *stream); +void LaunchTRSQRT_f32_merge_axis_16x72(void *a, void *b, void *stream); void LaunchTRSQRT_f32_32x32(void *a, void *b, void *stream); void LaunchTRSQRT_f16_16x64(void *a, void *b, void *stream); void LaunchTRSQRT_f16_32x32(void *a, void *b, void *stream); @@ -44,6 +45,7 @@ static const TestCase kCases[] = { {"f32_32x32", LaunchTRSQRT_f32_32x32, 32, 32, 32, 32, sizeof(float)}, {"f16_16x64", LaunchTRSQRT_f16_16x64, 16, 64, 16, 64, sizeof(uint16_t)}, {"f16_32x32", LaunchTRSQRT_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", LaunchTRSQRT_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/trsqrt.pto b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/trsqrt.pto index 7881bb3eab..93f5e4a6ae 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/trsqrt.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/trsqrt/trsqrt.pto @@ -178,4 +178,49 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TRSQRT_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.trsqrt ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshl/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tshl/cases.py index 4bc3084005..f7897ccda7 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshl/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshl/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 0, }, + { + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshl/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshl/launch.cpp index d58d324a45..c5b74d3ff7 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshl/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshl/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TSHL_i32_32x32(__gm__ int32_t *a, __gm__ int32 void LaunchTSHL_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream) { TSHL_i32_32x32<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TSHL_i32_merge_axis_16x72(__gm__ int32_t *a, __gm__ int32_t *b, __gm__ int32_t *c); + +void LaunchTSHL_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream) { + TSHL_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshl/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshl/main.cpp index cb35c3a31d..a13cadcfe6 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshl/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshl/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSHL_i32_16x64(int32_t *a, int32_t *b, int32_t *c, void *stream); +void LaunchTSHL_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream); + void LaunchTSHL_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream); using LaunchFn = void (*)(int32_t *, int32_t *, int32_t *, void *); @@ -40,6 +42,7 @@ struct TestCase { static const TestCase kCases[] = { {"i32_16x64", LaunchTSHL_i32_16x64, 16, 64, 16, 64, sizeof(int32_t)}, {"i32_32x32", LaunchTSHL_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, + {"i32_merge_axis_16x72", LaunchTSHL_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -142,4 +145,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshl/tshl.pto b/test/tilelang_st/npu/a5/src/st/testcase/tshl/tshl.pto index c002895d97..e004544616 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshl/tshl.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshl/tshl.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: i32 full-axis 16x72 (1152 elements) + func.func @TSHL_i32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%b : !pto.tile_buf) + + pto.tshl ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshls/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tshls/cases.py index 18cc991784..5199832bf2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshls/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshls/cases.py @@ -39,4 +39,11 @@ "valid_shape": (15, 192), "eps": 0, }, -] +{ + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshls/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshls/launch.cpp index 5e73430712..33a517f719 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshls/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshls/launch.cpp @@ -42,3 +42,10 @@ extern "C" __global__ AICORE void TSHLS_i16_15x192(__gm__ int16_t *src, __gm__ i void LaunchTSHLS_i16_15x192(int16_t *src, int16_t *dst, void *stream) { TSHLS_i16_15x192<<<1, nullptr, stream>>>((__gm__ int16_t *)src, (__gm__ int16_t *)dst, TSHLS_SCALAR); } + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TSHLS_i32_merge_axis_16x72(__gm__ int32_t *src, __gm__ int32_t *dst, int16_t scalar); + +void LaunchTSHLS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream) { + TSHLS_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, TSHLS_SCALAR); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshls/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshls/main.cpp index ca0e4d73e8..de10368815 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshls/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshls/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSHLS_i32_32x64(int32_t *src, int32_t *dst, void *stream); +void LaunchTSHLS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream); void LaunchTSHLS_i16_63x64(int16_t *src, int16_t *dst, void *stream); void LaunchTSHLS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTSHLS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -42,6 +43,7 @@ static const TestCase kCases[] = { {"i16_63x64", (void (*)(void*,void*,void*))LaunchTSHLS_i16_63x64, 63, 64, 63, 64, sizeof(int16_t)}, {"i32_31x128", (void (*)(void*,void*,void*))LaunchTSHLS_i32_31x128, 31, 128, 31, 128, sizeof(int32_t)}, {"i16_15x192", (void (*)(void*,void*,void*))LaunchTSHLS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, + {"i32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTSHLS_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshls/tshls.pto b/test/tilelang_st/npu/a5/src/st/testcase/tshls/tshls.pto index 6bd6e4fd39..1463eb368d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshls/tshls.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshls/tshls.pto @@ -173,4 +173,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: i16) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + pto.tshls ins(%src, %scalar : !pto.tile_buf, i16) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshr/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tshr/cases.py index 36075525bf..b1718126ca 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshr/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshr/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 0, }, + { + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshr/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshr/launch.cpp index 1d4f9cd1d6..3caeefe0f8 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshr/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshr/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TSHR_i32_32x32(__gm__ int32_t *a, __gm__ int32 void LaunchTSHR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream) { TSHR_i32_32x32<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TSHR_i32_merge_axis_16x72(__gm__ int32_t *a, __gm__ int32_t *b, __gm__ int32_t *c); + +void LaunchTSHR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream) { + TSHR_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshr/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshr/main.cpp index e99634eaec..38fd6e621d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshr/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshr/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSHR_i32_16x64(int32_t *a, int32_t *b, int32_t *c, void *stream); +void LaunchTSHR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream); + void LaunchTSHR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream); using LaunchFn = void (*)(int32_t *, int32_t *, int32_t *, void *); @@ -40,6 +42,7 @@ struct TestCase { static const TestCase kCases[] = { {"i32_16x64", LaunchTSHR_i32_16x64, 16, 64, 16, 64, sizeof(int32_t)}, {"i32_32x32", LaunchTSHR_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, + {"i32_merge_axis_16x72", LaunchTSHR_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -142,4 +145,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshr/tshr.pto b/test/tilelang_st/npu/a5/src/st/testcase/tshr/tshr.pto index ab6a9f36fb..59aede3da0 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshr/tshr.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshr/tshr.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: i32 full-axis 16x72 (1152 elements) + func.func @TSHR_i32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%b : !pto.tile_buf) + + pto.tshr ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/cases.py index 18cc991784..5199832bf2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/cases.py @@ -39,4 +39,11 @@ "valid_shape": (15, 192), "eps": 0, }, -] +{ + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/launch.cpp index e80c03e166..0cf3e1f628 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/launch.cpp @@ -42,3 +42,10 @@ extern "C" __global__ AICORE void TSHRS_i16_15x192(__gm__ int16_t *src, __gm__ i void LaunchTSHRS_i16_15x192(int16_t *src, int16_t *dst, void *stream) { TSHRS_i16_15x192<<<1, nullptr, stream>>>((__gm__ int16_t *)src, (__gm__ int16_t *)dst, TSHRS_SCALAR); } + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TSHRS_i32_merge_axis_16x72(__gm__ int32_t *src, __gm__ int32_t *dst, int16_t scalar); + +void LaunchTSHRS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream) { + TSHRS_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, TSHRS_SCALAR); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/main.cpp index 3afd710f9a..5087027586 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSHRS_i32_32x64(int32_t *src, int32_t *dst, void *stream); +void LaunchTSHRS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream); void LaunchTSHRS_i16_63x64(int16_t *src, int16_t *dst, void *stream); void LaunchTSHRS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTSHRS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -42,6 +43,7 @@ static const TestCase kCases[] = { {"i16_63x64", (void (*)(void*,void*,void*))LaunchTSHRS_i16_63x64, 63, 64, 63, 64, sizeof(int16_t)}, {"i32_31x128", (void (*)(void*,void*,void*))LaunchTSHRS_i32_31x128, 31, 128, 31, 128, sizeof(int32_t)}, {"i16_15x192", (void (*)(void*,void*,void*))LaunchTSHRS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, + {"i32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTSHRS_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/tshrs.pto b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/tshrs.pto index 94fe2c4f46..7ca59a686e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tshrs/tshrs.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tshrs/tshrs.pto @@ -173,4 +173,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: i16) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + pto.tshrs ins(%src, %scalar : !pto.tile_buf, i16) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/cases.py index 24fb10786b..550d0b3c97 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/cases.py @@ -72,4 +72,12 @@ "eps": 1e-7, "high_precision": True, }, +{ + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + "high_precision": False, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/launch.cpp index cb1d3cb3db..61b35e1f58 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/launch.cpp @@ -52,4 +52,11 @@ extern "C" __global__ AICORE void TSQRT_f16_64x64_hp2(__gm__ uint16_t *a, __gm__ void LaunchTSQRT_f16_64x64_hp2(void *a, void *b, void *stream) { TSQRT_f16_64x64_hp2<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TSQRT_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b); + +void LaunchTSQRT_f32_merge_axis_16x72(void *a, void *b, void *stream) { + TSQRT_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/main.cpp index d7be5f315d..d247bb4aee 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSQRT_f32_16x64(void *a, void *b, void *stream); +void LaunchTSQRT_f32_merge_axis_16x72(void *a, void *b, void *stream); void LaunchTSQRT_f32_32x32(void *a, void *b, void *stream); void LaunchTSQRT_f16_16x64(void *a, void *b, void *stream); void LaunchTSQRT_f16_32x32(void *a, void *b, void *stream); @@ -48,6 +49,7 @@ static const TestCase kCases[] = { {"f16_32x32", LaunchTSQRT_f16_32x32, 32, 32, 32, 32, sizeof(uint16_t)}, {"f32_64x64_hp1", LaunchTSQRT_f32_64x64_hp1, 64, 64, 64, 64, sizeof(float)}, {"f16_64x64_hp2", LaunchTSQRT_f16_64x64_hp2, 64, 64, 64, 64, sizeof(uint16_t)}, + {"f32_merge_axis_16x72", LaunchTSQRT_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/tsqrt.pto b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/tsqrt.pto index f7e49af9f5..61d12b01d4 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/tsqrt.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsqrt/tsqrt.pto @@ -261,4 +261,49 @@ pto.tstore ins(%b : !pto.tile_buf) outs(%b_part : !pto.partition_tensor_view<1x1x1x64x64xf16>) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TSQRT_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + + pto.tsqrt ins(%a : !pto.tile_buf) + outs(%b : !pto.tile_buf) + + pto.tstore ins(%b : !pto.tile_buf) + outs(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsub/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tsub/cases.py index b71da2e9b7..cb277978fd 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsub/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsub/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 1e-6, }, + { + "name": "f32_merge_axis_16x72", + "dtype": np.float32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 1e-6, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsub/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsub/launch.cpp index 256c0ed07f..3e8c12c074 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsub/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsub/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TSUB_f32_32x32(__gm__ float *a, __gm__ float * void LaunchTSUB_f32_32x32(float *a, float *b, float *c, void *stream) { TSUB_f32_32x32<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); -} \ No newline at end of file +} + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TSUB_f32_merge_axis_16x72(__gm__ float *a, __gm__ float *b, __gm__ float *c); + +void LaunchTSUB_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream) { + TSUB_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)a, (__gm__ float *)b, (__gm__ float *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsub/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsub/main.cpp index b5e338d4ba..9297de89bf 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsub/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsub/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSUB_f32_16x64(float *a, float *b, float *c, void *stream); +void LaunchTSUB_f32_merge_axis_16x72(float *a, float *b, float *c, void *stream); void LaunchTSUB_f32_32x32(float *a, float *b, float *c, void *stream); using LaunchFn = void (*)(float *, float *, float *, void *); @@ -40,6 +41,7 @@ struct TestCase { static const TestCase kCases[] = { {"f32_16x64", LaunchTSUB_f32_16x64, 16, 64, 16, 64, sizeof(float)}, {"f32_32x32", LaunchTSUB_f32_32x32, 32, 32, 32, 32, sizeof(float)}, + {"f32_merge_axis_16x72", LaunchTSUB_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsub/tsub.pto b/test/tilelang_st/npu/a5/src/st/testcase/tsub/tsub.pto index 3246f3bc10..29d358b12d 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsub/tsub.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsub/tsub.pto @@ -120,4 +120,62 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: f32 full-axis 16x72 (1152 elements) + func.func @TSUB_f32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%b : !pto.tile_buf) + + pto.tsub ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/cases.py index af6b6b4254..0afdf87c49 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/cases.py @@ -19,4 +19,5 @@ {"name": "i16_15x192", "dtype": np.int16, "shape": (15, 192), "valid_shape": (15, 192), "eps": 0}, {"name": "f32_7x448", "dtype": np.float32, "shape": (7, 448), "valid_shape": (7, 448), "eps": 1e-6}, {"name": "f32_256x16", "dtype": np.float32, "shape": (256, 16), "valid_shape": (256, 16), "eps": 1e-6}, -] +{"name": "f32_merge_axis_16x72", "dtype": np.float32, "shape": (16, 72), "valid_shape": (16, 72), "eps": 1e-6}, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/launch.cpp index d511cf09dc..01a05314f2 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/launch.cpp @@ -56,3 +56,10 @@ extern "C" __global__ AICORE void TSUBS_f32_256x16(__gm__ float *src, __gm__ flo void LaunchTSUBS_f32_256x16(float *src, float *dst, void *stream) { TSUBS_f32_256x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TSUBS_SCALAR_F32); } + +// Merge-axis case: f32 merge axis 16x72 +extern "C" __global__ AICORE void TSUBS_f32_merge_axis_16x72(__gm__ float *src, __gm__ float *dst, float scalar); + +void LaunchTSUBS_f32_merge_axis_16x72(float *src, float *dst, void *stream) { + TSUBS_f32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, TSUBS_SCALAR_F32); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/main.cpp index 40509f5786..669357de99 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTSUBS_f32_32x64(float *src, float *dst, void *stream); +void LaunchTSUBS_f32_merge_axis_16x72(float *src, float *dst, void *stream); void LaunchTSUBS_f16_63x64(uint16_t *src, uint16_t *dst, void *stream); void LaunchTSUBS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTSUBS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -46,6 +47,7 @@ static const TestCase kCases[] = { {"i16_15x192", (void (*)(void*,void*,void*))LaunchTSUBS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, {"f32_7x448", (void (*)(void*,void*,void*))LaunchTSUBS_f32_7x448, 7, 448, 7, 448, sizeof(float)}, {"f32_256x16", (void (*)(void*,void*,void*))LaunchTSUBS_f32_256x16, 256, 16, 256, 16, sizeof(float)}, + {"f32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTSUBS_f32_merge_axis_16x72, 16, 72, 16, 72, sizeof(float)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/tsubs.pto b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/tsubs.pto index 4e9dcaef18..a63536ac92 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/tsubs/tsubs.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/tsubs/tsubs.pto @@ -253,4 +253,47 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: f32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xf32> -> !pto.partition_tensor_view<1x1x1x16x72xf32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + outs(%src : !pto.tile_buf) + pto.tsubs ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xf32>) + return + } } diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txor/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/txor/cases.py index c710ea612b..d79b79e871 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txor/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/txor/cases.py @@ -37,4 +37,11 @@ "valid_shape": (32, 32), "eps": 0, }, + { + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, ] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txor/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/txor/launch.cpp index 90fd20459c..a4af8cac4e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txor/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/txor/launch.cpp @@ -24,4 +24,11 @@ extern "C" __global__ AICORE void TXOR_i32_32x32(__gm__ int32_t *a, __gm__ int32 void LaunchTXOR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream) { TXOR_i32_32x32<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); -} \ No newline at end of file +} + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TXOR_i32_merge_axis_16x72(__gm__ int32_t *a, __gm__ int32_t *b, __gm__ int32_t *c); + +void LaunchTXOR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream) { + TXOR_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)a, (__gm__ int32_t *)b, (__gm__ int32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txor/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/txor/main.cpp index 838ff0de1c..1af4561389 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txor/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/txor/main.cpp @@ -23,6 +23,8 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTXOR_i32_16x64(int32_t *a, int32_t *b, int32_t *c, void *stream); +void LaunchTXOR_i32_merge_axis_16x72(int32_t *a, int32_t *b, int32_t *c, void *stream); + void LaunchTXOR_i32_32x32(int32_t *a, int32_t *b, int32_t *c, void *stream); using LaunchFn = void (*)(int32_t *, int32_t *, int32_t *, void *); @@ -40,6 +42,7 @@ struct TestCase { static const TestCase kCases[] = { {"i32_16x64", LaunchTXOR_i32_16x64, 16, 64, 16, 64, sizeof(int32_t)}, {"i32_32x32", LaunchTXOR_i32_32x32, 32, 32, 32, 32, sizeof(int32_t)}, + {"i32_merge_axis_16x72", LaunchTXOR_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); @@ -142,4 +145,4 @@ int main(int argc, char *argv[]) { aclFinalize(); return rc; -} \ No newline at end of file +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txor/txor.pto b/test/tilelang_st/npu/a5/src/st/testcase/txor/txor.pto index 41de710fa5..cc9271fcf1 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txor/txor.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/txor/txor.pto @@ -122,4 +122,63 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) return } -} \ No newline at end of file + + + + + // Merge-axis case: i32 full-axis 16x72 (1152 elements) + func.func @TXOR_i32_merge_axis_16x72(%a_ptr: !pto.ptr, %b_ptr: !pto.ptr, %c_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c1024 = arith.constant 1152 : index + + %a_view = pto.make_tensor_view %a_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %b_view = pto.make_tensor_view %b_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %c_view = pto.make_tensor_view %c_ptr, + shape = [%c1, %c1, %c1, %c16, %c64], + strides = [%c1024, %c1024, %c1024, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %a_part = pto.partition_view %a_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %b_part = pto.partition_view %b_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %c_part = pto.partition_view %c_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %c = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%a_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%a : !pto.tile_buf) + pto.tload ins(%b_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%b : !pto.tile_buf) + + pto.txor ins(%a, %b, %c : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs(%c : !pto.tile_buf) + + pto.tstore ins(%c : !pto.tile_buf) + outs(%c_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txors/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/txors/cases.py index 9b652056d9..4d2d7c3cba 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txors/cases.py +++ b/test/tilelang_st/npu/a5/src/st/testcase/txors/cases.py @@ -45,4 +45,11 @@ "valid_shape": (15, 192), "eps": 0, }, -] +{ + "name": "i32_merge_axis_16x72", + "dtype": np.int32, + "shape": (16, 72), + "valid_shape": (16, 72), + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txors/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/txors/launch.cpp index f61619d9f1..a0c9e56956 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txors/launch.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/txors/launch.cpp @@ -39,3 +39,10 @@ extern "C" __global__ AICORE void TXORS_i16_15x192(__gm__ int16_t *src, __gm__ i void LaunchTXORS_i16_15x192(int16_t *src, int16_t *dst, void *stream) { TXORS_i16_15x192<<<1, nullptr, stream>>>((__gm__ int16_t *)src, (__gm__ int16_t *)dst, (int16_t)3); } + +// Merge-axis case: i32 merge axis 16x72 +extern "C" __global__ AICORE void TXORS_i32_merge_axis_16x72(__gm__ int32_t *src, __gm__ int32_t *dst, int32_t scalar); + +void LaunchTXORS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream) { + TXORS_i32_merge_axis_16x72<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, (int32_t)3); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txors/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/txors/main.cpp index f46282f01a..77a1b0af3b 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txors/main.cpp +++ b/test/tilelang_st/npu/a5/src/st/testcase/txors/main.cpp @@ -23,6 +23,7 @@ using namespace PtoTestCommon; // Kernel launch wrappers (defined in launch.cpp) void LaunchTXORS_i32_32x64(int32_t *src, int32_t *dst, void *stream); +void LaunchTXORS_i32_merge_axis_16x72(int32_t *src, int32_t *dst, void *stream); void LaunchTXORS_i16_63x64(int16_t *src, int16_t *dst, void *stream); void LaunchTXORS_i32_31x128(int32_t *src, int32_t *dst, void *stream); void LaunchTXORS_i16_15x192(int16_t *src, int16_t *dst, void *stream); @@ -42,6 +43,7 @@ static const TestCase kCases[] = { {"i16_63x64", (void (*)(void*,void*,void*))LaunchTXORS_i16_63x64, 63, 64, 63, 64, sizeof(int16_t)}, {"i32_31x128", (void (*)(void*,void*,void*))LaunchTXORS_i32_31x128, 31, 128, 31, 128, sizeof(int32_t)}, {"i16_15x192", (void (*)(void*,void*,void*))LaunchTXORS_i16_15x192, 15, 192, 15, 192, sizeof(int16_t)}, + {"i32_merge_axis_16x72", (void (*)(void*,void*,void*))LaunchTXORS_i32_merge_axis_16x72, 16, 72, 16, 72, sizeof(int32_t)}, }; static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); diff --git a/test/tilelang_st/npu/a5/src/st/testcase/txors/txors.pto b/test/tilelang_st/npu/a5/src/st/testcase/txors/txors.pto index 17ec4041f8..8d7bef994e 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/txors/txors.pto +++ b/test/tilelang_st/npu/a5/src/st/testcase/txors/txors.pto @@ -186,4 +186,50 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, %dst_ptr: !pto.ptr, %scalar: i32) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 16 : index + %c64 = arith.constant 72 : index + %c2048 = arith.constant 1152 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c64], + strides = [%c2048, %c2048, %c2048, %c64, %c1] + : !pto.tensor_view<1x1x1x16x72xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c64] + : !pto.tensor_view<1x1x1x16x72xi32> -> !pto.partition_tensor_view<1x1x1x16x72xi32> + + %src = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + outs(%src : !pto.tile_buf) + pto.txors ins(%src, %scalar, %tmp : !pto.tile_buf, i32, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tstore ins(%dst : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x72xi32>) + return + } }