From 9246d2f24b416b0dd48306ecbcf2c51a00c7b6eb Mon Sep 17 00:00:00 2001 From: caojian5 Date: Tue, 2 Jun 2026 20:44:50 +0800 Subject: [PATCH 1/2] add tinsert op & st --- lib/TileOps/tinsert_template.py | 1459 +++++++++++++++++ .../npu/a5/src/st/testcase/CMakeLists.txt | 2 + .../a5/src/st/testcase/tinsert/CMakeLists.txt | 9 + .../npu/a5/src/st/testcase/tinsert/cases.py | 87 + .../npu/a5/src/st/testcase/tinsert/compare.py | 66 + .../a5/src/st/testcase/tinsert/gen_data.py | 35 + .../npu/a5/src/st/testcase/tinsert/launch.cpp | 56 + .../npu/a5/src/st/testcase/tinsert/main.cpp | 150 ++ .../a5/src/st/testcase/tinsert/tinsert.pto | 595 +++++++ .../st/testcase/tinsert_vec/CMakeLists.txt | 9 + .../a5/src/st/testcase/tinsert_vec/cases.py | 73 + .../a5/src/st/testcase/tinsert_vec/compare.py | 66 + .../src/st/testcase/tinsert_vec/gen_data.py | 35 + .../a5/src/st/testcase/tinsert_vec/launch.cpp | 44 + .../a5/src/st/testcase/tinsert_vec/main.cpp | 158 ++ .../st/testcase/tinsert_vec/tinsert_vec.pto | 305 ++++ 16 files changed, 3149 insertions(+) create mode 100644 lib/TileOps/tinsert_template.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp create mode 100644 test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto diff --git a/lib/TileOps/tinsert_template.py b/lib/TileOps/tinsert_template.py new file mode 100644 index 0000000000..6ce4f714b0 --- /dev/null +++ b/lib/TileOps/tinsert_template.py @@ -0,0 +1,1459 @@ +# 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. + +"""TileLang DSL template for pto.tinsert - tile data insertion + +Implements six data movement paths using @pto.ckernel: + - Acc→Mat NZ: pto.mte_l0c_l1 (copy_matrix_cc_to_cbuf) + - Acc→Vec ND/DN/NZ: pto.mte_l0c_ub (copy_matrix_cc_to_ub) + - Vec→Vec ND/NZ: pto.copy_ubuf_to_ubuf + - Vec→Mat NZ/ND: pto.mte_ub_l1 (copy_ubuf_to_cbuf) +""" + +import tilelang_dsl as pto + +BLOCK_BYTE_SIZE = 32 +BLOCK_BYTE_BITS = 256 +FRACTAL_NZ_ROW = 16 + + +# --------------------------------------------------------------------------- +# Constraint functions +# --------------------------------------------------------------------------- + + +def _acc_to_mat_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "mat" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _acc_to_vec_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + ) + + +def _acc_to_vec_dn_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + ) + + +def _acc_to_vec_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_vec_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + and not (src.valid_shape[0] == 1 and src.valid_shape[1] == 1) + ) + + +def _vec_to_vec_nd_scalar_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + and src.valid_shape[0] == 1 + and src.valid_shape[1] == 1 + ) + + +def _vec_to_vec_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.COL_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_mat_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.COL_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR + and dst.memory_space == "mat" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_mat_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "mat" + and dst.config.s_layout == pto.SLayout.NONE_BOX + and src.shape[1] * pto.bytewidth(pto.ScalarType(src.dtype)) >= BLOCK_BYTE_SIZE + ) + + +# --------------------------------------------------------------------------- +# Acc -> Mat +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.i8, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + constraints=[_acc_to_mat_constraint], +) +def template_tinsert_acc_to_mat( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + n_size = (valid_cols + c0_size - 1) // c0_size * c0_size + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = src.shape[0] * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + constraints=[_acc_to_mat_constraint], +) +def template_tinsert_acc_to_mat_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + n_size = (valid_cols + c0_size - 1) // c0_size * c0_size + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = src.shape[0] * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (ND, NONE_BOX) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_nd_constraint], +) +def template_tinsert_acc_to_vec_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size + + dst_cols = dst.shape[1] + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_cols * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + priority=10, + constraints=[_acc_to_vec_nd_constraint], +) +def template_tinsert_acc_to_vec_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size + + dst_cols = dst.shape[1] + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_cols * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (DN, NONE_BOX) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_dn_constraint], +) +def template_tinsert_acc_to_vec_dn( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows_raw = src.shape[0] + valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size + valid_cols = src.shape[1] + + dst_rows = dst.shape[0] + dst_offset = index_col * dst_rows + index_row + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.i32), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + priority=10, + constraints=[_acc_to_vec_dn_constraint], +) +def template_tinsert_acc_to_vec_dn_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows_raw = src.shape[0] + valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size + valid_cols = src.shape[1] + + dst_rows = dst.shape[0] + dst_offset = index_col * dst_rows + index_row + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (NZ, ROW_MAJOR) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_nz_constraint], +) +def template_tinsert_acc_to_vec_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + + if pto.constexpr(dst_elem == pto.f32): + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == BLOCK_BYTE_BITS): + valid_cols_align = c0_size + else: + valid_cols_align = FRACTAL_NZ_ROW + else: + valid_cols_align = c0_size + valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.f16), + ], + priority=10, + constraints=[_acc_to_vec_nz_constraint], +) +def template_tinsert_acc_to_vec_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + + if pto.constexpr(dst_elem == pto.f32): + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == BLOCK_BYTE_BITS): + valid_cols_align = c0_size + else: + valid_cols_align = FRACTAL_NZ_ROW + else: + valid_cols_align = c0_size + valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + ) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Vec (ND, ROW_MAJOR + NONE_BOX) - O1 +# --------------------------------------------------------------------------- + + +_VEC_TO_VEC_DTYPES = [ + (pto.f16, pto.i64, pto.i64, pto.f16), + (pto.bf16, pto.i64, pto.i64, pto.bf16), + (pto.f32, pto.i64, pto.i64, pto.f32), + (pto.i32, pto.i64, pto.i64, pto.i32), + (pto.i8, pto.i64, pto.i64, pto.i8), +] + +_VEC_TO_VEC_BASIC_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + +_VEC_TO_MAT_SPLIT_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + lanes = pto.get_lanes(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_stride = src.shape[1] + dst_stride = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_stride + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + src_stride_bytes = src_stride * elem_bytes + dst_stride_bytes = dst_stride * elem_bytes + strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 + + if pto.constexpr(strides_aligned): + if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: + if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + else: + repeat_times = (valid_cols + lanes - 1) // lanes + for i in range(0, valid_rows, 1): + remained = valid_cols + for j in range(0, repeat_times, 1): + pred, remained = pto.make_mask(dtype, remained) + src_off = i * src_stride + j * lanes + dst_off = i * dst_stride + j * lanes + data = pto.vlds(pto.addptr(src_ptr, src_off), 0) + pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + lanes = pto.get_lanes(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_stride = src.shape[1] + dst_stride = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_stride + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + src_stride_bytes = src_stride * elem_bytes + dst_stride_bytes = dst_stride * elem_bytes + strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 + + if pto.constexpr(strides_aligned): + if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: + if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + else: + repeat_times = (valid_cols + lanes - 1) // lanes + for i in range(0, valid_rows, 1): + remained = valid_cols + for j in range(0, repeat_times, 1): + pred, remained = pto.make_mask(dtype, remained) + src_off = i * src_stride + j * lanes + dst_off = i * dst_stride + j * lanes + data = pto.vlds(pto.addptr(src_ptr, src_off), 0) + pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nd_scalar_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_scalar( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dst_stride = dst.shape[1] + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + src_val = pto.load_scalar(src_ptr, 0) + dst_elem_offset = index_row * dst_stride + index_col + pto.store_scalar(src_val, dst_ptr, dst_elem_offset) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nd_scalar_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_scalar_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dst_stride = dst.shape[1] + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + src_val = pto.load_scalar(src_ptr, 0) + dst_elem_offset = index_row * dst_stride + index_col + pto.store_scalar(src_val, dst_ptr, dst_elem_offset) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Vec (NZ, COL_MAJOR + ROW_MAJOR) - O2 +# --------------------------------------------------------------------------- + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (NZ, COL_MAJOR + ROW_MAJOR) - O3 +# --------------------------------------------------------------------------- + + +_VEC_TO_MAT_DTYPES = [ + (pto.f16, pto.i64, pto.i64, pto.f16), + (pto.bf16, pto.i64, pto.i64, pto.bf16), + (pto.f32, pto.i64, pto.i64, pto.f32), + (pto.i32, pto.i64, pto.i64, pto.i32), + (pto.i8, pto.i64, pto.i64, pto.i8), +] + +_VEC_TO_MAT_BASIC_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_DTYPES, + constraints=[_vec_to_mat_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_BASIC_DTYPES, + constraints=[_vec_to_mat_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (ND, ROW_MAJOR + NONE_BOX) - O4 +# --------------------------------------------------------------------------- + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_DTYPES, + constraints=[_vec_to_mat_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_cols = src.shape[1] + dst_cols = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + + if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_BASIC_DTYPES, + constraints=[_vec_to_mat_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_cols = src.shape[1] + dst_cols = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + + if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (NZ, Split2/Split4) - O5 +# Split large-tile NZ DMA into 2/4 independent segments to reduce L1 bank +# conflicts (mirrors pto-isa TInsertMode::SPLIT2 / SPLIT4). +# --------------------------------------------------------------------------- + + +def _make_split_template(split_count): + @pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_SPLIT_DTYPES, + name=f"template_tinsert_vec_to_mat_nz_split{split_count}", + constraints=[_vec_to_mat_nz_constraint], + advanced=True, + ) + def _split_fn(src: pto.Tile, index_row: pto.i32, index_col: pto.i32, dst: pto.Tile): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + aligned_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_base = pto.addptr(dst.as_ptr(), dst_offset) + + total_burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = aligned_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = aligned_rows + 1 + else: + src_stride_rows = aligned_rows + src_gap = src_stride_rows - aligned_rows + dst_gap = dst_rows - aligned_rows + + part_num = total_burst_num // split_count + last_num = total_burst_num - part_num * (split_count - 1) + src_block_size = (burst_len + src_gap) * BLOCK_BYTE_SIZE // elem_bytes + dst_block_size = dst_rows * c0_size + + pto.mte_ub_l1(src_ptr, dst_base, burst_len, nburst=(part_num, src_gap, dst_gap)) + + src_ptr1 = pto.addptr(src_ptr, part_num * src_block_size) + dst_ptr1 = pto.addptr(dst_base, part_num * dst_block_size) + if pto.constexpr(split_count == 2): + pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(last_num, src_gap, dst_gap)) + else: + pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(part_num, src_gap, dst_gap)) + + if pto.constexpr(split_count == 4): + src_ptr2 = pto.addptr(src_ptr, 2 * part_num * src_block_size) + dst_ptr2 = pto.addptr(dst_base, 2 * part_num * dst_block_size) + pto.mte_ub_l1(src_ptr2, dst_ptr2, burst_len, nburst=(part_num, src_gap, dst_gap)) + + src_ptr3 = pto.addptr(src_ptr, 3 * part_num * src_block_size) + dst_ptr3 = pto.addptr(dst_base, 3 * part_num * dst_block_size) + pto.mte_ub_l1(src_ptr3, dst_ptr3, burst_len, nburst=(last_num, src_gap, dst_gap)) + return None + + return _split_fn + + +template_tinsert_vec_to_mat_nz_split2 = _make_split_template(2) +template_tinsert_vec_to_mat_nz_split4 = _make_split_template(4) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt index 22fbc3f731..e2be53efbe 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt @@ -110,6 +110,8 @@ endfunction() # -------------------------------------------------------------------------- set(ALL_TESTCASES tadd + tinsert + tinsert_vec tsub tmul tdiv diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt new file mode 100644 index 0000000000..6878082ab1 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt @@ -0,0 +1,9 @@ +# 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. + +pto_tilelang_cube_st(tinsert PTO_LEVEL level3) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py new file mode 100644 index 0000000000..53da565802 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py @@ -0,0 +1,87 @@ +# 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. + +# coding=utf-8 + +"""Test cases for pto.tinsert ST (Acc->Mat NZ, Acc->Vec ND/NZ).""" + +import numpy as np + + +CASES = [ + { + "name": "acc2mat_f16_16x16", + "kernel": "TINSERT_acc2mat_f16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_f16_32x32", + "kernel": "TINSERT_acc2mat_f16_32x32", + "m": 32, "k": 32, "n": 32, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_bf16_16x16", + "kernel": "TINSERT_acc2mat_bf16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_f32_16x16", + "kernel": "TINSERT_acc2mat_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2vec_nd_f16_16x16", + "kernel": "TINSERT_acc2vec_nd_f16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nd", + "has_output": True, + "eps": 1e-2, + }, + { + "name": "acc2vec_nd_f32_16x16", + "kernel": "TINSERT_acc2vec_nd_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nd", + "has_output": True, + "eps": 1e-2, + }, + { + "name": "acc2vec_nz_f32_16x16", + "kernel": "TINSERT_acc2vec_nz_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nz", + "has_output": True, + "eps": 1e-2, + }, +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py new file mode 100644 index 0000000000..90ae4a1667 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py @@ -0,0 +1,66 @@ +# 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. + +# coding=utf-8 + +import os +import sys +import numpy as np + +from cases import CASES +from st_common import result_cmp, style_fail, style_pass + + +def main(): + case_filter = sys.argv[1] if len(sys.argv) > 1 else None + + all_passed = True + for case in CASES: + if case_filter is not None and case["name"] != case_filter: + continue + + case_dir = case["name"] + m, n = case["m"], case["n"] + dtype_out = case["dtype_out"] + + if not case["has_output"]: + print(style_pass(f"[INFO] {case['name']}: compile-only (no output comparison)")) + continue + + golden_path = os.path.join(case_dir, "golden.bin") + output_path = os.path.join(case_dir, "output.bin") + + if not os.path.exists(output_path): + print(style_fail(f"[ERROR] {case['name']}: output.bin not found (kernel may not have written output)")) + all_passed = False + continue + + golden = np.fromfile(golden_path, dtype=dtype_out) + output = np.fromfile(output_path, dtype=dtype_out) + + if golden.shape != output.shape: + print(style_fail( + f"[ERROR] {case['name']}: shape mismatch golden={golden.shape} output={output.shape}" + )) + all_passed = False + continue + + ok = result_cmp(golden, output, case["eps"]) + if ok: + print(style_pass(f"[INFO] {case['name']}: compare passed")) + else: + print(style_fail(f"[ERROR] {case['name']}: compare failed")) + all_passed = False + + if not all_passed: + sys.exit(2) + print(style_pass("[INFO] all cases passed")) + + +if __name__ == "__main__": + main() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py new file mode 100644 index 0000000000..ddceacec0e --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py @@ -0,0 +1,35 @@ +# 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. + +# coding=utf-8 + +import os +import numpy as np + +from cases import CASES +from st_common import setup_case_rng, save_case_data + + +for case in CASES: + setup_case_rng(case) + m, k, n = case["m"], case["k"], case["n"] + dtype = case["dtype"] + dtype_out = case["dtype_out"] + + A = np.random.uniform(-1.0, 1.0, size=(m, k)).astype(dtype) + B = np.random.uniform(-1.0, 1.0, size=(k, n)).astype(dtype) + golden_f32 = np.matmul(A.astype(np.float32), B.astype(np.float32)) + golden = golden_f32.astype(dtype_out) + + data = {"input1": A, "input2": B, "golden": golden} + + save_case_data(case["name"], data) + print( + f"[INFO] gen_data: {case['name']} A=({m},{k}) B=({k},{n}) " + f"dtype={dtype.__name__} dtype_out={dtype_out.__name__}" + ) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp new file mode 100644 index 0000000000..9671721364 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp @@ -0,0 +1,56 @@ +// 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. + +#include + +#ifndef AICORE +#define AICORE [aicore] +#endif + +extern "C" __global__ AICORE void TINSERT_acc2mat_f16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_f16_32x32( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_bf16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nd_f16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nd_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nz_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); + +void LaunchAcc2Mat_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_f16_32x32(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_f16_32x32<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_bf16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_bf16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_f32_16x16(uint16_t *a, uint16_t *b, uint32_t *c, void *stream) { + TINSERT_acc2mat_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecND_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nd_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecND_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nd_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecNZ_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nz_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp new file mode 100644 index 0000000000..7da0c77cb4 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp @@ -0,0 +1,150 @@ +// 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. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include +#include +#include + +using namespace PtoTestCommon; + +using LaunchFn = void (*)(uint16_t *, uint16_t *, uint16_t *, void *); + +void LaunchAcc2Mat_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_f16_32x32(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_bf16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_f32_16x16(uint16_t *a, uint16_t *b, uint32_t *c, void *stream); +void LaunchAcc2VecND_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2VecND_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2VecNZ_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t m, k, n; + bool has_output; + size_t out_elem_bytes; +}; + +static const TestCase kCases[] = { + {"acc2mat_f16_16x16", LaunchAcc2Mat_f16_16x16, 16, 16, 16, false, 2}, + {"acc2mat_f16_32x32", LaunchAcc2Mat_f16_32x32, 32, 32, 32, false, 2}, + {"acc2mat_bf16_16x16", LaunchAcc2Mat_bf16_16x16, 16, 16, 16, false, 2}, + {"acc2mat_f32_16x16", reinterpret_cast(LaunchAcc2Mat_f32_16x16), 16, 16, 16, false, 4}, + {"acc2vec_nd_f16_16x16", LaunchAcc2VecND_f16_16x16, 16, 16, 16, true, 4}, + {"acc2vec_nd_f32_16x16", LaunchAcc2VecND_f32_16x16, 16, 16, 16, true, 4}, + {"acc2vec_nz_f32_16x16", LaunchAcc2VecNZ_f32_16x16, 16, 16, 16, true, 4}, +}; +static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); + +static int RunCase(const TestCase &tc, int deviceId, aclrtStream stream) { + (void)deviceId; + int rc = 0; + const size_t aElems = tc.m * tc.k; + const size_t bElems = tc.k * tc.n; + const size_t outElems = tc.m * tc.n; + const size_t aBytes = aElems * sizeof(uint16_t); + const size_t bBytes = bElems * sizeof(uint16_t); + const size_t outBytes = outElems * tc.out_elem_bytes; + size_t aFileSize = aBytes; + size_t bFileSize = bBytes; + + std::printf("[INFO] === case: %s (m=%zu, k=%zu, n=%zu) ===\n", tc.name, tc.m, tc.k, tc.n); + + std::string caseDir = std::string("./") + tc.name; + + void *aHost = nullptr, *bHost = nullptr, *outHost = nullptr; + void *aDev = nullptr, *bDev = nullptr, *outDev = nullptr; + + aclrtMallocHost(&aHost, aBytes); + aclrtMallocHost(&bHost, bBytes); + aclrtMallocHost(&outHost, outBytes); + aclrtMalloc(&aDev, aBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&bDev, bBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDev, outBytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input1.bin").c_str(), aFileSize, aHost, aBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input1.bin\n", caseDir.c_str()); + rc = 1; + } + if (rc == 0 && !ReadFile((caseDir + "/input2.bin").c_str(), bFileSize, bHost, bBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input2.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(aDev, aBytes, aHost, aBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(bDev, bBytes, bHost, bBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemset(outDev, outBytes, 0, outBytes); + + tc.launch( + static_cast(aDev), + static_cast(bDev), + static_cast(outDev), + stream + ); + + aclrtSynchronizeStream(stream); + + if (tc.has_output) { + aclrtMemcpy(outHost, outBytes, outDev, outBytes, ACL_MEMCPY_DEVICE_TO_HOST); + if (!WriteFile((caseDir + "/output.bin").c_str(), outHost, outBytes)) { + std::fprintf(stderr, "[ERROR] failed to write %s/output.bin\n", caseDir.c_str()); + rc = 1; + } + } + } + + if (aDev) aclrtFree(aDev); + if (bDev) aclrtFree(bDev); + if (outDev) aclrtFree(outDev); + if (aHost) aclrtFreeHost(aHost); + if (bHost) aclrtFreeHost(bHost); + if (outHost) aclrtFreeHost(outHost); + + if (rc == 0) + std::printf("[INFO] case %s done\n", tc.name); + return rc; +} + +int main(int argc, char *argv[]) { + const char *caseFilter = (argc > 1) ? argv[1] : nullptr; + + int rc = 0; + int deviceId = 0; + aclrtStream stream = nullptr; + + aclInit(nullptr); + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) { + deviceId = std::atoi(envDevice); + } + aclrtSetDevice(deviceId); + aclrtCreateStream(&stream); + + for (size_t i = 0; i < kNumCases; ++i) { + if (caseFilter != nullptr && std::strcmp(kCases[i].name, caseFilter) != 0) { + continue; + } + int ret = RunCase(kCases[i], deviceId, stream); + if (ret != 0) { + std::fprintf(stderr, "[ERROR] case %s failed\n", kCases[i].name); + rc = 1; + break; + } + } + + if (stream != nullptr) + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); + + return rc; +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto new file mode 100644 index 0000000000..5385679f23 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto @@ -0,0 +1,595 @@ +// 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 and copyright holder. + +// ST kernels for pto.tinsert ckernel templates. +// Each kernel: load A,B from GM (MTE) -> matmul -> tinsert(Acc->dst) -> store result to GM. +// Uses low-level MTE ops (level3) for boundary data movement. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + // ------------------------------------------------------------------------- + // Case 3: Acc → Vec ND (16×16×16, f16×f16→f32→f16) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→vec_ND) → mte(l0c→GM) + // Output: L0C accumulator (f32, ND layout via nz2nd). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nd_f16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %dst_vec = pto.tile_buf_addr %dst_vec_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 5: Acc → Vec NZ (16×16×16, f16×f16→f32→f32) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→vec_NZ) → mte(l0c→GM) + // Output: L0C accumulator (f32, ND layout via nz2nd). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nz_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c64_i64 = arith.constant 64 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %dst_vec = pto.tile_buf_addr %dst_vec_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 1: Acc → Mat NZ (16×16×16, f16×f16→f32→f16) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→mat_NZ) + // Note: compile+run only (no output comparison). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 2: Acc → Mat NZ (32×32×32, f16×f16→f32→f16) — larger tile variant + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f16_32x32(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c32_i64, %c32_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c32_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c32_i64, %c32_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c32_i64, %c32_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c32_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c32_i64, %c32_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Mat NZ (16×16×16, f16×f16→f32→bf16) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_bf16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Mat NZ (16×16×16, f16×f16→f32→f32) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Vec ND (16×16×16, f16×f16→f32→f32) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nd_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt new file mode 100644 index 0000000000..be1519c4ae --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt @@ -0,0 +1,9 @@ +# 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. + +pto_tilelang_vec_st(tinsert_vec) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py new file mode 100644 index 0000000000..4f2c5d1b66 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.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. + +# coding=utf-8 + +"""Test cases for pto.tinsert Vec->Vec ND vkernel ST.""" + +import numpy as np +from ml_dtypes import bfloat16 + + +CASES = [ + { + "name": "vec2vec_nd_f16_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00", + "dtype": np.float16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_f16_16x16_into_32x32_idx816", + "kernel": "TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816", + "dtype": np.float16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 8, + "index_col": 16, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_f32_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00", + "dtype": np.float32, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-6, + }, + { + "name": "vec2vec_nd_bf16_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00", + "dtype": bfloat16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_i32_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00", + "dtype": np.int32, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py new file mode 100644 index 0000000000..0272300355 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py @@ -0,0 +1,66 @@ +# 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. + +# coding=utf-8 + +import os +import sys +import numpy as np + +from cases import CASES +from st_common import result_cmp, style_fail, style_pass + + +def main(): + case_filter = sys.argv[1] if len(sys.argv) > 1 else None + + all_passed = True + for case in CASES: + if case_filter is not None and case["name"] != case_filter: + continue + + case_dir = case["name"] + dtype = case["dtype"] + dst_rows, dst_cols = case["dst_shape"] + + if not case["has_output"]: + print(style_pass(f"[INFO] {case['name']}: compile-only (no output comparison)")) + continue + + golden_path = os.path.join(case_dir, "golden.bin") + output_path = os.path.join(case_dir, "output.bin") + + if not os.path.exists(output_path): + print(style_fail(f"[ERROR] {case['name']}: output.bin not found")) + all_passed = False + continue + + golden = np.fromfile(golden_path, dtype=dtype).reshape(dst_rows, dst_cols) + output = np.fromfile(output_path, dtype=dtype).reshape(dst_rows, dst_cols) + + if golden.shape != output.shape: + print(style_fail( + f"[ERROR] {case['name']}: shape mismatch golden={golden.shape} output={output.shape}" + )) + all_passed = False + continue + + ok = result_cmp(golden, output, case["eps"]) + if ok: + print(style_pass(f"[INFO] {case['name']}: compare passed")) + else: + print(style_fail(f"[ERROR] {case['name']}: compare failed")) + all_passed = False + + if not all_passed: + sys.exit(2) + print(style_pass("[INFO] all cases passed")) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py new file mode 100644 index 0000000000..e3faf1b586 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py @@ -0,0 +1,35 @@ +# 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. + +# coding=utf-8 + +import numpy as np + +from cases import CASES +from st_common import setup_case_rng, save_case_data + + +for case in CASES: + setup_case_rng(case) + dtype = case["dtype"] + src_rows, src_cols = case["src_shape"] + dst_rows, dst_cols = case["dst_shape"] + idx_row, idx_col = case["index_row"], case["index_col"] + + src = np.random.uniform(-1.0, 1.0, size=(src_rows, src_cols)).astype(dtype) + dst = np.random.uniform(-1.0, 1.0, size=(dst_rows, dst_cols)).astype(dtype) + + golden = dst.copy() + golden[idx_row:idx_row + src_rows, idx_col:idx_col + src_cols] = src + + data = {"input1": src, "input2": dst, "golden": golden} + save_case_data(case["name"], data) + print( + f"[INFO] gen_data: {case['name']} src=({src_rows},{src_cols}) dst=({dst_rows},{dst_cols}) " + f"idx=({idx_row},{idx_col}) dtype={dtype.__name__}" + ) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp new file mode 100644 index 0000000000..4e36cf8bf1 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp @@ -0,0 +1,44 @@ +// 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. + +#include + +#ifndef AICORE +#define AICORE [aicore] +#endif + +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00( + __gm__ float *src, __gm__ float *dst, __gm__ float *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00( + __gm__ int32_t *src, __gm__ int32_t *dst, __gm__ int32_t *out); + +void LaunchVec2VecND_f16_16x16_into_32x32_idx00(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_f16_16x16_into_32x32_idx816(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_f32_16x16_into_32x32_idx00(float *src, float *dst, float *out, void *stream) { + TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, (__gm__ float *)out); +} + +void LaunchVec2VecND_bf16_16x16_into_32x32_idx00(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_i32_16x16_into_32x32_idx00(int32_t *src, int32_t *dst, int32_t *out, void *stream) { + TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, (__gm__ int32_t *)out); +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp new file mode 100644 index 0000000000..0dd7b743b5 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp @@ -0,0 +1,158 @@ +// 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. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include +#include +#include + +using namespace PtoTestCommon; + +void LaunchVec2VecND_f16_16x16_into_32x32_idx00( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_f16_16x16_into_32x32_idx816( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_f32_16x16_into_32x32_idx00( + float *src, float *dst, float *out, void *stream); +void LaunchVec2VecND_bf16_16x16_into_32x32_idx00( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_i32_16x16_into_32x32_idx00( + int32_t *src, int32_t *dst, int32_t *out, void *stream); + +using LaunchFn = void (*)(void *, void *, void *, void *); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t src_rows, src_cols; + size_t dst_rows, dst_cols; + size_t idx_row, idx_col; + bool has_output; + size_t elem_bytes; +}; + +static const TestCase kCases[] = { + {"vec2vec_nd_f16_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_f16_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 2}, + {"vec2vec_nd_f16_16x16_into_32x32_idx816", + reinterpret_cast(LaunchVec2VecND_f16_16x16_into_32x32_idx816), + 16, 16, 32, 32, 8, 16, true, 2}, + {"vec2vec_nd_f32_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_f32_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 4}, + {"vec2vec_nd_bf16_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_bf16_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 2}, + {"vec2vec_nd_i32_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_i32_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 4}, +}; +static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); + +static int RunCase(const TestCase &tc, int deviceId, aclrtStream stream) { + (void)deviceId; + int rc = 0; + const size_t srcElems = tc.src_rows * tc.src_cols; + const size_t dstElems = tc.dst_rows * tc.dst_cols; + const size_t srcBytes = srcElems * tc.elem_bytes; + const size_t dstBytes = dstElems * tc.elem_bytes; + size_t srcFileSize = srcBytes; + size_t dstFileSize = dstBytes; + + std::printf("[INFO] === case: %s src=%zux%zu dst=%zux%zu idx=(%zu,%zu) ===\n", + tc.name, tc.src_rows, tc.src_cols, tc.dst_rows, tc.dst_cols, + tc.idx_row, tc.idx_col); + + std::string caseDir = std::string("./") + tc.name; + + void *srcHost = nullptr, *dstHost = nullptr, *outHost = nullptr; + void *srcDev = nullptr, *dstDev = nullptr, *outDev = nullptr; + + aclrtMallocHost(&srcHost, srcBytes); + aclrtMallocHost(&dstHost, dstBytes); + aclrtMallocHost(&outHost, dstBytes); + aclrtMalloc(&srcDev, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&dstDev, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDev, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input1.bin").c_str(), srcFileSize, srcHost, srcBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input1.bin\n", caseDir.c_str()); + rc = 1; + } + if (rc == 0 && !ReadFile((caseDir + "/input2.bin").c_str(), dstFileSize, dstHost, dstBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input2.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(srcDev, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(dstDev, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemset(outDev, dstBytes, 0, dstBytes); + + tc.launch(srcDev, dstDev, outDev, stream); + + aclrtSynchronizeStream(stream); + + if (tc.has_output) { + aclrtMemcpy(outHost, dstBytes, outDev, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST); + if (!WriteFile((caseDir + "/output.bin").c_str(), outHost, dstBytes)) { + std::fprintf(stderr, "[ERROR] failed to write %s/output.bin\n", caseDir.c_str()); + rc = 1; + } + } + } + + if (srcDev) aclrtFree(srcDev); + if (dstDev) aclrtFree(dstDev); + if (outDev) aclrtFree(outDev); + if (srcHost) aclrtFreeHost(srcHost); + if (dstHost) aclrtFreeHost(dstHost); + if (outHost) aclrtFreeHost(outHost); + + if (rc == 0) + std::printf("[INFO] case %s done\n", tc.name); + return rc; +} + +int main(int argc, char *argv[]) { + const char *caseFilter = (argc > 1) ? argv[1] : nullptr; + + int rc = 0; + int deviceId = 0; + aclrtStream stream = nullptr; + + aclInit(nullptr); + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) { + deviceId = std::atoi(envDevice); + } + aclrtSetDevice(deviceId); + aclrtCreateStream(&stream); + + for (size_t i = 0; i < kNumCases; ++i) { + if (caseFilter != nullptr && std::strcmp(kCases[i].name, caseFilter) != 0) { + continue; + } + int ret = RunCase(kCases[i], deviceId, stream); + if (ret != 0) { + std::fprintf(stderr, "[ERROR] case %s failed\n", kCases[i].name); + rc = 1; + break; + } + } + + if (stream != nullptr) + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); + + return rc; +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto new file mode 100644 index 0000000000..29d023c19f --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto @@ -0,0 +1,305 @@ +// 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 OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software and copyright holder. + +// ST kernels for pto.tinsert Vec->Vec ND vkernel templates. +// Pattern: tload(src->src_tile) + tload(dst->dst_tile) + tinsert(src_tile,idx->dst_tile) + tstore(dst_tile->out) + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f16 -> dst=(32,32) f16, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f16 -> dst=(32,32) f16, index=(8,16) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx8 = arith.constant 8 : index + %idx16 = arith.constant 16 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx8, %idx16 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f32 -> dst=(32,32) f32, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf32> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf32> -> !pto.partition_tensor_view<1x1x1x16x16xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf32> -> !pto.partition_tensor_view<1x1x1x32x32xf32> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf32> -> !pto.partition_tensor_view<1x1x1x32x32xf32> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf32>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf32>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf32>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) bf16 -> dst=(32,32) bf16, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xbf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xbf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xbf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xbf16> -> !pto.partition_tensor_view<1x1x1x16x16xbf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xbf16> -> !pto.partition_tensor_view<1x1x1x32x32xbf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xbf16> -> !pto.partition_tensor_view<1x1x1x32x32xbf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xbf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xbf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xbf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) i32 -> dst=(32,32) i32, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xi32> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xi32> -> !pto.partition_tensor_view<1x1x1x16x16xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xi32> -> !pto.partition_tensor_view<1x1x1x32x32xi32> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xi32> -> !pto.partition_tensor_view<1x1x1x32x32xi32> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xi32>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xi32>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xi32>) + return + } + +} From 23aa40d55aaca8c82726009aa2057d05e1f772f9 Mon Sep 17 00:00:00 2001 From: caojian5 Date: Wed, 3 Jun 2026 09:47:59 +0800 Subject: [PATCH 2/2] =?UTF-8?q?add=20tinsert=20op=20&=20st=20|=20added=20f?= =?UTF-8?q?ramework-level=20support=20for=20the=20acc=E2=86=92vec?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- lib/PTO/IR/PTO.cpp | 22 ++++++++++++++++++- lib/PTO/IR/VPTO.cpp | 17 ++++++++++---- .../lit/pto/tinsert_verify_a5_invalid_loc.pto | 2 +- 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6952a7e288..72bee29582 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -5324,9 +5324,29 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { return success(); } + // A5 acc->vec path (L0C->UB, ND/NZ modes in pto-isa). + if (*srcSpace == pto::AddressSpace::ACC && *dstSpace == pto::AddressSpace::VEC) { + if (!isColMajorRowMajorNZ(srcTb)) + return emitOpError( + "expects A5 acc->vec tinsert src to use blayout=col_major and slayout=row_major"); + bool dstIsND = isRowMajorNoneBoxND(dstTb); + bool dstIsNZ = isColMajorRowMajorNZ(dstTb); + if (!dstIsND && !dstIsNZ) + return emitOpError( + "expects A5 acc->vec tinsert dst to use ND (row_major/none_box) or NZ (col_major/row_major) layout"); + bool okTypes = (srcElem.isF32() && + (dstElem.isF16() || dstElem.isBF16() || dstElem.isF32())) || + (srcElem.isInteger(32) && dstElem.isInteger(32)); + if (!okTypes) + return emitOpError( + "expects A5 acc->vec tinsert element types to be " + "(src=f32,dst=f16/bf16/f32) or (src=i32,dst=i32)"); + return success(); + } + return emitOpError( "expects A5 tinsert to use a supported src/dst loc pair: " - "acc->mat, vec->mat, or vec->vec"); + "acc->mat, acc->vec, vec->mat, or vec->vec"); }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index afe98d76bb..e897066eaa 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -2625,12 +2625,21 @@ static ParseResult parseStructuredAccStoreClauses( if (seenClause) { if (failed(parser.parseOptionalComma())) return success(); + } else { + // First iteration: optionally consume the comma between operands and + // clauses. For some callers (MteL0cUbOp after dst_mode parsing with + // parseOptionalComma already consumed), the comma may already be gone. + (void)parser.parseOptionalComma(); } StringRef keyword; - if (parser.parseKeyword(&keyword)) { + if (failed(parser.parseOptionalKeyword(&keyword))) { + // No keyword found. If this is the first iteration (no clauses yet), + // that means there are zero clauses — return success without error. + // On subsequent iterations, a keyword is expected after the comma. if (!seenClause) return success(); - return failure(); + return parser.emitError(parser.getCurrentLocation(), + "expected mte_l0c clause keyword"); } seenClause = true; @@ -6996,7 +7005,7 @@ ParseResult MteL0cL1Op::parse(OpAsmParser &parser, OperationState &result) { parseRequiredOperandWithComma(parser, m) || parseRequiredOperandWithComma(parser, n) || parseRequiredOperandWithComma(parser, srcStride) || - parseRequiredOperandWithComma(parser, dstStride) || + parser.parseOperand(dstStride) || parseStructuredAccStoreClauses(parser, state) || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) return failure(); @@ -7164,7 +7173,7 @@ ParseResult MteL0cGmOp::parse(OpAsmParser &parser, OperationState &result) { parseRequiredOperandWithComma(parser, srcStride) || parseRequiredOperandWithComma(parser, dstStride) || parseRequiredOperandWithComma(parser, sid) || - parseRequiredOperandWithComma(parser, l2CacheCtrl) || + parser.parseOperand(l2CacheCtrl) || parseStructuredAccStoreClauses(parser, state) || parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) return failure(); diff --git a/test/lit/pto/tinsert_verify_a5_invalid_loc.pto b/test/lit/pto/tinsert_verify_a5_invalid_loc.pto index 7125132e5a..e01a1c18e9 100644 --- a/test/lit/pto/tinsert_verify_a5_invalid_loc.pto +++ b/test/lit/pto/tinsert_verify_a5_invalid_loc.pto @@ -13,4 +13,4 @@ module attributes {"pto.target_arch" = "a5"} { } } -// CHECK: error: 'pto.tinsert' op expects A5 tinsert to use a supported src/dst loc pair: acc->mat, vec->mat, or vec->vec +// CHECK: error: 'pto.tinsert' op expects A5 tinsert to use a supported src/dst loc pair: acc->mat, acc->vec, vec->mat, or vec->vec