diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6952a7e288..84cd7d5fbe 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -5166,15 +5166,8 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { return emitOpError("expects A5 textract src to use a supported mat blayout/slayout combination"); if (*dstSpace == pto::AddressSpace::LEFT || *dstSpace == pto::AddressSpace::RIGHT) { - auto indexRowVal = getConstantIntegerValueEx( - getIndexRow(), /*includeIndexAndIntOpsInConstFold=*/true); - auto indexColVal = getConstantIntegerValueEx( - getIndexCol(), /*includeIndexAndIntOpsInConstFold=*/true); - if (!indexRowVal || !indexColVal || *indexRowVal != 0 || - *indexColVal != 0) - return emitOpError( - "expects A5 mat->left/right textract to have indexRow=0 and " - "indexCol=0 (offset extraction not yet supported)"); + // Non-zero indexRow/indexCol is supported via start_row/start_col + // on mte_l1_l0a/l0b (PR #469). } if (*dstSpace == pto::AddressSpace::LEFT) { if (dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::ColMajor) || diff --git a/lib/TileOps/textract_template.py b/lib/TileOps/textract_template.py new file mode 100644 index 0000000000..108bb5878f --- /dev/null +++ b/lib/TileOps/textract_template.py @@ -0,0 +1,425 @@ +# 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.textract and pto.textract_fp (A5). + +TEXTRACT extracts a sub-tile window from src into dst at offset (indexRow, indexCol): + dst[i, j] = src[i + indexRow, j + indexCol] + +Supported data-flow directions (A5): + +Cube path (Mat -> Left/Right): + The mte_l1_l0a/l0b transpose flag is set based on src tile layout: + - Same fractal (src.col_major blayout + src.row_major slayout): + transpose=False — direct extraction (TExtractToA / TExtractToACompact) + - Cross fractal (src.row_major blayout + src.col_major slayout): + transpose=True — transposed extraction (TExtractToATransCompact / TExtractToA) + Offset extraction (indexRow/indexCol != 0) is supported via start_row/start_col + keywords on mte_l1_l0a/l0b (PR #469). + +Fix-pipe path (Acc -> Mat with FP quantization): + - textract_fp : PIPE_FIX, uses pto.mte_l0c_l1 with pre_quant keyword + Currently constrained to indexRow=0, indexCol=0. The index_row/index_col + parameters are accepted but unused in the lowering because the start + position for fix-pipe copy is implicit (whole-tile extraction from L0C). + Future work: pass index offset to mte_l0c_l1 when PTO-ISA supports it. + +Vector path (Vec -> Vec): + - Vec -> Vec ND : PIPE_V, uses pto.vlds/vsts subscript syntax + Note: Uses vector load/store instead of DMA (copy_ubuf_to_ubuf) because + textract requires arbitrary offset support (index_row/index_col), which + DMA block copy cannot provide. + +Blocked paths: + - Mat -> ScaleLeft/ScaleRight : mte_l1_l0a_mx/l0b_mx reject memory_space="scaling" dst + - Vec -> Mat : now available via pto.mte_ub_l1 (PR #405) +""" + +import tilelang_dsl as pto + +_FP_QUANT_MODE_MAP = { + (pto.f32, pto.si8): "qf322b8_pre_vec", + (pto.f32, pto.ui8): "qf322b8_pre_vec", + (pto.f32, pto.f16): "qf322f16_pre_vec", + (pto.f32, pto.bf16): "qf322bf16_pre_vec", + (pto.f32, pto.f32): "qf322f32_pre_vec", + (pto.si32, pto.si8): "req8_vec", + (pto.si32, pto.ui8): "req8_vec", + (pto.si32, pto.f16): "deqf16_vec", + (pto.si32, pto.bf16): "qs322bf16_pre_vec", +} + + +def _is_same_fractal_as_left(src) -> bool: + if src.config is None: + return True + return (src.config.b_layout != pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR) + + +def _is_cross_fractal_as_left(src) -> bool: + if src.config is None: + return False + return (src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.COL_MAJOR) + + +def _textract_cube_dst_is_left_same_fractal(src, index_row, index_col, dst) -> bool: + dst_ms = dst.memory_space + if not (dst_ms == "left" if isinstance(dst_ms, str) + else dst_ms.value == "left"): + return False + if not _is_same_fractal_as_left(src): + return False + return True + + +def _textract_cube_dst_is_left_cross_fractal(src, index_row, index_col, dst) -> bool: + dst_ms = dst.memory_space + if not (dst_ms == "left" if isinstance(dst_ms, str) + else dst_ms.value == "left"): + return False + if not _is_cross_fractal_as_left(src): + return False + return True + + +def _is_same_fractal_as_right(src) -> bool: + if src.config is None: + return True + return (src.config.b_layout != pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR) + + +def _is_cross_fractal_as_right(src) -> bool: + if src.config is None: + return False + return (src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.COL_MAJOR) + + +def _textract_cube_dst_is_right_same_fractal(src, index_row, index_col, dst) -> bool: + dst_ms = dst.memory_space + if not (dst_ms == "right" if isinstance(dst_ms, str) + else dst_ms.value == "right"): + return False + if not _is_same_fractal_as_right(src): + return False + return True + + +def _textract_cube_dst_is_right_cross_fractal(src, index_row, index_col, dst) -> bool: + dst_ms = dst.memory_space + if not (dst_ms == "right" if isinstance(dst_ms, str) + else dst_ms.value == "right"): + return False + if not _is_cross_fractal_as_right(src): + return False + return True + + +def _textract_vec2vec_nd_constraint(src, index_row, index_col, dst) -> bool: + src_ms = src.memory_space + dst_ms = dst.memory_space + src_is_ub = (src_ms == "ub" if isinstance(src_ms, str) + else src_ms.value == "ub") + dst_is_ub = (dst_ms == "ub" if isinstance(dst_ms, str) + else dst_ms.value == "ub") + if not (src_is_ub and dst_is_ub): + return False + if src.config is None or dst.config is None: + return False + if src.config.b_layout != pto.BLayout.ROW_MAJOR: + return False + if src.config.s_layout != pto.SLayout.NONE_BOX: + return False + if dst.config.b_layout != pto.BLayout.ROW_MAJOR: + return False + if dst.config.s_layout != pto.SLayout.NONE_BOX: + return False + if src.dtype != dst.dtype: + return False + return True + + +def _textract_fp_acc2mat_constraint(src, fp, index_row, index_col, dst) -> bool: + src_ms = src.memory_space + fp_ms = fp.memory_space + dst_ms = dst.memory_space + if not (src_ms == "acc" if isinstance(src_ms, str) + else src_ms.value == "acc"): + return False + if not (fp_ms == "scaling" if isinstance(fp_ms, str) + else fp_ms.value == "scaling"): + return False + if not (dst_ms == "mat" if isinstance(dst_ms, str) + else dst_ms.value == "mat"): + return False + index_row_val = index_row.value if hasattr(index_row, 'value') else None + index_col_val = index_col.value if hasattr(index_col, 'value') else None + if index_row_val is not None and index_row_val != 0: + return False + if index_col_val is not None and index_col_val != 0: + return False + return (src.dtype, dst.dtype) in _FP_QUANT_MODE_MAP + + +def _make_fp_constraint(src_dtype, dst_dtype): + def _fp_constraint(src, fp, index_row, index_col, dst) -> bool: + src_ms = src.memory_space + fp_ms = fp.memory_space + dst_ms = dst.memory_space + if not (src_ms == "acc" if isinstance(src_ms, str) + else src_ms.value == "acc"): + return False + if not (fp_ms == "scaling" if isinstance(fp_ms, str) + else fp_ms.value == "scaling"): + return False + if not (dst_ms == "mat" if isinstance(dst_ms, str) + else dst_ms.value == "mat"): + return False + index_row_val = index_row.value if hasattr(index_row, 'value') else None + index_col_val = index_col.value if hasattr(index_col, 'value') else None + if index_row_val is not None and index_row_val != 0: + return False + if index_col_val is not None and index_col_val != 0: + return False + return src.dtype == src_dtype and dst.dtype == dst_dtype + return _fp_constraint + + +@pto.ckernel( + target="a5", + op="pto.textract", + constraints=[_textract_cube_dst_is_left_same_fractal], +) +def template_textract_mat2left(src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, k = dst.valid_shape + pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), m, k, + start_row=index_row, start_col=index_col) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract", + constraints=[_textract_cube_dst_is_left_cross_fractal], +) +def template_textract_mat2left_trans(src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, k = dst.valid_shape + pto.mte_l1_l0a(src.as_ptr(), dst.as_ptr(), m, k, + start_row=index_row, start_col=index_col, transpose=True) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract", + constraints=[_textract_cube_dst_is_right_same_fractal], +) +def template_textract_mat2right(src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + k, n = dst.valid_shape + pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), k, n, + start_row=index_row, start_col=index_col) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract", + constraints=[_textract_cube_dst_is_right_cross_fractal], +) +def template_textract_mat2right_trans(src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + k, n = dst.valid_shape + pto.mte_l1_l0b(src.as_ptr(), dst.as_ptr(), k, n, + start_row=index_row, start_col=index_col, transpose=True) + return None + + +@pto.vkernel( + target="a5", + op="pto.textract", + advanced=True, + constraints=[_textract_vec2vec_nd_constraint], +) +def template_textract_vec2vec_nd(src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + dtype = dst.element_type + lanes = pto.get_lanes(dtype) + valid_rows, valid_cols = dst.valid_shape + + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + mask, remained = pto.make_mask(dtype, remained) + data = pto.vlds(src[index_row + row, index_col + col:]) + pto.vsts(data, dst[row, col:], mask) + + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.si8),), + constraints=[_make_fp_constraint(pto.f32, pto.si8)], +) +def template_textract_fp_f32_si8(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qf322b8_pre_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.ui8),), + constraints=[_make_fp_constraint(pto.f32, pto.ui8)], +) +def template_textract_fp_f32_ui8(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qf322b8_pre_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.f16),), + constraints=[_make_fp_constraint(pto.f32, pto.f16)], +) +def template_textract_fp_f32_f16(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qf322f16_pre_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.bf16),), + constraints=[_make_fp_constraint(pto.f32, pto.bf16)], +) +def template_textract_fp_f32_bf16(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qf322bf16_pre_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.f32, pto.f32, pto.i32, pto.i32, pto.f32),), + constraints=[_make_fp_constraint(pto.f32, pto.f32)], +) +def template_textract_fp_f32_f32(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qf322f32_pre_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.si8),), + constraints=[_make_fp_constraint(pto.si32, pto.si8)], +) +def template_textract_fp_si32_si8(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "req8_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.ui8),), + constraints=[_make_fp_constraint(pto.si32, pto.ui8)], +) +def template_textract_fp_si32_ui8(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "req8_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.f16),), + constraints=[_make_fp_constraint(pto.si32, pto.f16)], +) +def template_textract_fp_si32_f16(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "deqf16_vec")) + return None + + +@pto.ckernel( + target="a5", + op="pto.textract_fp", + dtypes=((pto.si32, pto.f32, pto.i32, pto.i32, pto.bf16),), + constraints=[_make_fp_constraint(pto.si32, pto.bf16)], +) +def template_textract_fp_si32_bf16(src: pto.Tile, fp: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile): + m, n = dst.valid_shape + src_stride = src.shape[0] + dst_stride = dst.shape[0] + pto.mte_l0c_l1(src.as_ptr(), dst.as_ptr(), m, n, src_stride, dst_stride, + pre_quant=(fp.as_ptr(), "qs322bf16_pre_vec")) + return None \ No newline at end of file diff --git a/test/lit/pto/textract_verify_a5_mat_lr_nonzero_offset.pto b/test/lit/pto/textract_verify_a5_mat_lr_nonzero_offset.pto index 81662b25dc..99e434c40e 100644 --- a/test/lit/pto/textract_verify_a5_mat_lr_nonzero_offset.pto +++ b/test/lit/pto/textract_verify_a5_mat_lr_nonzero_offset.pto @@ -1,15 +1,38 @@ -// RUN: not ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s +// 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 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. +// RUN: ptoas --pto-arch=a5 %s -o /dev/null 2>&1 | FileCheck %s --allow-empty + +// Verify that non-zero offset extraction is now accepted at PTO IR level +// (previously rejected, now supported via PR #469 start_row/start_col). + +// CHECK-NOT: error + + module attributes {"pto.target_arch" = "a5"} { + func.func @textract_a5_mat_left_nonzero_offset() { + %c0 = arith.constant 0 : index + %c4 = arith.constant 4 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract ins(%src, %c4, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + } -// CHECK: error: 'pto.textract' op expects A5 mat->left/right textract to have indexRow=0 and indexCol=0 (offset extraction not yet supported) diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto new file mode 100644 index 0000000000..20ee4c110b --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_acc2mat.pto @@ -0,0 +1,64 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_FP_A2M + +// CHECK-NOT: pto.textract_fp + +// CHECK: pto.set_fpc + +// CHECK: pto.copy_matrix_cc_to_cbuf + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_A2M() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_all_dtypes.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_all_dtypes.pto new file mode 100644 index 0000000000..6c81b662b8 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_fp_all_dtypes.pto @@ -0,0 +1,138 @@ +// 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. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + +// Verify all 9 FP dtype pairs expand into set_fpc + copy_matrix_cc_to_cbuf +// (no remaining textract_fp op). + +// CHECK-DAG: func.func @TEXTRACT_FP_F32_F32 +// CHECK-DAG: func.func @TEXTRACT_FP_F32_BF16 +// CHECK-DAG: func.func @TEXTRACT_FP_F32_F16 +// CHECK-DAG: func.func @TEXTRACT_FP_F32_SI8 +// CHECK-DAG: func.func @TEXTRACT_FP_F32_UI8 +// CHECK-DAG: func.func @TEXTRACT_FP_SI32_F16 +// CHECK-DAG: func.func @TEXTRACT_FP_SI32_BF16 +// CHECK-DAG: func.func @TEXTRACT_FP_SI32_SI8 +// CHECK-DAG: func.func @TEXTRACT_FP_SI32_UI8 +// CHECK-DAG: func.func @TEXTRACT_FP_A2M +// CHECK-NOT: pto.textract_fp + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_F32_F32() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_F32_BF16() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_F32_F16() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_F32_SI8() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_F32_UI8() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_SI32_F16() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_SI32_BF16() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_SI32_SI8() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_SI32_UI8() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func @TEXTRACT_FP_A2M() attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + +} \ No newline at end of file diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left.pto new file mode 100644 index 0000000000..79f493ee17 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2L + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_ca + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_bf16.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_bf16.pto new file mode 100644 index 0000000000..d96e8683c1 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_bf16.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2L_BF16 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_ca + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L_BF16() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_cm_rm.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_cm_rm.pto new file mode 100644 index 0000000000..591f3169b8 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_cm_rm.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2L_CM_RM + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_ca + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L_CM_RM() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_f32.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_f32.pto new file mode 100644 index 0000000000..598adf9e99 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2left_f32.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2L_F32 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_ca + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L_F32() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right.pto new file mode 100644 index 0000000000..76fcbfd726 --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2R + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_cb + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2R() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_bf16.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_bf16.pto new file mode 100644 index 0000000000..0baff9031b --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_bf16.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2R_BF16 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_cb + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2R_BF16() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_cm_rm.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_cm_rm.pto new file mode 100644 index 0000000000..9f9ea6e41c --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_cm_rm.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2R_CM_RM + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_cb + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2R_CM_RM() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_f32.pto b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_f32.pto new file mode 100644 index 0000000000..e6e078aaff --- /dev/null +++ b/test/lit/vpto/cube/expand_tile_op_tilelang_textract_mat2right_f32.pto @@ -0,0 +1,52 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_M2R_F32 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.load_cbuf_to_cb + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2R_F32() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto new file mode 100644 index 0000000000..787f69be8b --- /dev/null +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_dst_loc.pto @@ -0,0 +1,58 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects dst to use loc=mat + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_INVALID_DST_LOC() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_dtype_pair.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_dtype_pair.pto new file mode 100644 index 0000000000..d30dd4537d --- /dev/null +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_dtype_pair.pto @@ -0,0 +1,58 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract_fp element types to be + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_INVALID_DTYPE_PAIR() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_fp_loc.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_fp_loc.pto new file mode 100644 index 0000000000..1c663e26c2 --- /dev/null +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_fp_loc.pto @@ -0,0 +1,58 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects fp to use loc=scaling + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_INVALID_FP_LOC() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_src_layout.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_src_layout.pto new file mode 100644 index 0000000000..72c1c72ccf --- /dev/null +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_src_layout.pto @@ -0,0 +1,58 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects src to use blayout=col_major and slayout=row_major + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_INVALID_SRC_LAYOUT() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_fp_verify_invalid_src_loc.pto b/test/lit/vpto/cube/textract_fp_verify_invalid_src_loc.pto new file mode 100644 index 0000000000..6238851405 --- /dev/null +++ b/test/lit/vpto/cube/textract_fp_verify_invalid_src_loc.pto @@ -0,0 +1,58 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects src to use loc=acc + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_INVALID_SRC_LOC() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %fp = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract_fp ins(%src, %fp, %c0, %c0 : !pto.tile_buf, + + !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_dtype_mismatch.pto b/test/lit/vpto/cube/textract_verify_invalid_dtype_mismatch.pto new file mode 100644 index 0000000000..cde3ebe18b --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_dtype_mismatch.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects src and dst to have the same element type + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_DTYPE_MISMATCH() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_left_layout.pto b/test/lit/vpto/cube/textract_verify_invalid_left_layout.pto new file mode 100644 index 0000000000..f38b388b77 --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_left_layout.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract src to use a supported mat blayout/slayout combination + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L_SRC_LAYOUT_INVALID() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_loc_pair.pto b/test/lit/vpto/cube/textract_verify_invalid_loc_pair.pto new file mode 100644 index 0000000000..2f54ba2bd8 --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_loc_pair.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract to use a supported src/dst loc pair + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_UNSUPPORTED_LOC_PAIR() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_offset.pto b/test/lit/vpto/cube/textract_verify_invalid_offset.pto new file mode 100644 index 0000000000..389efd1e0b --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_offset.pto @@ -0,0 +1,51 @@ +// 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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + +// Verify that non-zero offset extraction is supported via start_row/start_col. + +// CHECK-LABEL: func.func @TEXTRACT_M2L_OFFSET +// CHECK: pto.load_cbuf_to_ca {{.*}}, {{.*}}, %c4_i64, %c0_i64 + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2L_OFFSET() attributes {pto.aicore} { + + %c4 = arith.constant 4 : index + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c4, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_right_layout.pto b/test/lit/vpto/cube/textract_verify_invalid_right_layout.pto new file mode 100644 index 0000000000..783469ff88 --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_right_layout.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 right dst to use row_major blayout and col_major slayout + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_M2R_LAYOUT_INVALID() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_src_acc.pto b/test/lit/vpto/cube/textract_verify_invalid_src_acc.pto new file mode 100644 index 0000000000..3e288f4364 --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_src_acc.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract to use a supported src/dst loc pair + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_INVALID_SRC_ACC() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/cube/textract_verify_invalid_vec_left.pto b/test/lit/vpto/cube/textract_verify_invalid_vec_left.pto new file mode 100644 index 0000000000..096053ad69 --- /dev/null +++ b/test/lit/vpto/cube/textract_verify_invalid_vec_left.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract to use a supported src/dst loc pair + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_INVALID_VEC_LEFT() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd.pto b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd.pto new file mode 100644 index 0000000000..4a27b496e5 --- /dev/null +++ b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd.pto @@ -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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK: func.func @TEXTRACT_V2V_ND + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.vecscope + +// CHECK: pto.vlds + +// CHECK: pto.vsts + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_ND() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_bf16.pto b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_bf16.pto new file mode 100644 index 0000000000..ba3b1bc8cd --- /dev/null +++ b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_bf16.pto @@ -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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK: func.func @TEXTRACT_V2V_ND_BF16 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.vecscope + +// CHECK: pto.vlds + +// CHECK: pto.vsts + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_ND_BF16() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_f16.pto b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_f16.pto new file mode 100644 index 0000000000..0673b743f3 --- /dev/null +++ b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_f16.pto @@ -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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK: func.func @TEXTRACT_V2V_ND_F16 + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.vecscope + +// CHECK: pto.vlds + +// CHECK: pto.vsts + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_ND_F16() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_offset.pto b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_offset.pto new file mode 100644 index 0000000000..739106002f --- /dev/null +++ b/test/lit/vpto/expand_tile_op_tilelang_textract_vec2vec_nd_offset.pto @@ -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 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. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>/dev/null | FileCheck %s + + + +// CHECK-LABEL: func.func @TEXTRACT_V2V_ND_OFFSET + +// CHECK-NOT: pto.textract ins + +// CHECK: pto.vecscope + +// CHECK: pto.vlds + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_ND_OFFSET() attributes {pto.aicore} { + + %c8 = arith.constant 8 : index + + %c16 = arith.constant 16 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c8, %c16 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/textract_verify_invalid_vec2mat_loc_pair.pto b/test/lit/vpto/textract_verify_invalid_vec2mat_loc_pair.pto new file mode 100644 index 0000000000..647dda09e7 --- /dev/null +++ b/test/lit/vpto/textract_verify_invalid_vec2mat_loc_pair.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto --enable-tile-op-expand --emit-vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: found no registered kernel + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2M_NO_TEMPLATE() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/textract_verify_invalid_vec2vec_dtype_mismatch.pto b/test/lit/vpto/textract_verify_invalid_vec2vec_dtype_mismatch.pto new file mode 100644 index 0000000000..04176f17cf --- /dev/null +++ b/test/lit/vpto/textract_verify_invalid_vec2vec_dtype_mismatch.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects src and dst to have the same element type + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_DTYPE_MISMATCH() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/textract_verify_invalid_vec2vec_layout.pto b/test/lit/vpto/textract_verify_invalid_vec2vec_layout.pto new file mode 100644 index 0000000000..0327a51ffe --- /dev/null +++ b/test/lit/vpto/textract_verify_invalid_vec2vec_layout.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 vec->vec textract src/dst to use ND layout + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_LAYOUT_INVALID() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} diff --git a/test/lit/vpto/textract_verify_invalid_vec2vec_unsupported_dtype.pto b/test/lit/vpto/textract_verify_invalid_vec2vec_unsupported_dtype.pto new file mode 100644 index 0000000000..e682a2c27a --- /dev/null +++ b/test/lit/vpto/textract_verify_invalid_vec2vec_unsupported_dtype.pto @@ -0,0 +1,48 @@ +// 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 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. +// RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %s -o - 2>&1 | FileCheck %s + + + +// CHECK: expects A5 textract element type to be an fp8/f16/bf16/f32 or int8 family type + + + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_UNSUPPORTED_DTYPE() attributes {pto.aicore} { + + %c0 = arith.constant 0 : index + + %src = pto.alloc_tile + + : !pto.tile_buf + + %dst = pto.alloc_tile + + : !pto.tile_buf + + + + pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + + outs(%dst : !pto.tile_buf) + + return + + } + +} 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..9bc5ad33e5 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt @@ -194,6 +194,9 @@ set(ALL_TESTCASES tfmods tcmps tmatmul + textract + textract_fp + textract_v2v ) if((TEST_CASE IN_LIST ALL_TESTCASES) OR (TEST_CASE STREQUAL "all")) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/textract/CMakeLists.txt new file mode 100644 index 0000000000..76a55a9d6d --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/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(textract PTO_LEVEL level3) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/textract/cases.py new file mode 100644 index 0000000000..e29599365f --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/cases.py @@ -0,0 +1,33 @@ +# 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 + + +CASES = [ + { + "name": "mat2left_f16_16x16", + "dtype_src": np.float16, + "dtype_id": np.float16, + "shape_src": (16, 16), + "shape_id": (16, 16), + "shape_out": (16, 16), + "eps": 1e-2, + }, + { + "name": "mat2right_f16_16x16", + "dtype_src": np.float16, + "dtype_id": np.float16, + "shape_src": (16, 16), + "shape_id": (16, 16), + "shape_out": (16, 16), + "eps": 1e-2, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/textract/compare.py new file mode 100644 index 0000000000..0a4b11f0d5 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/compare.py @@ -0,0 +1,45 @@ +# 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"] + shape_out = case["shape_out"] + golden = np.fromfile(os.path.join(case_dir, "golden.bin"), dtype=np.float32).reshape(shape_out) + output = np.fromfile(os.path.join(case_dir, "output.bin"), dtype=np.float32).reshape(shape_out) + + 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/textract/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/textract/gen_data.py new file mode 100644 index 0000000000..5c2d5f8020 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/gen_data.py @@ -0,0 +1,36 @@ +# 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 + + +import os + + +for case in CASES: + setup_case_rng(case) + name = case["name"] + + if name.startswith("mat2left"): + src = np.random.uniform(-1.0, 1.0, size=case["shape_src"]).astype(case["dtype_src"]) + id_mat = np.eye(case["shape_id"][0], case["shape_id"][1], dtype=case["dtype_id"]) + golden = np.matmul(src.astype(np.float32), id_mat.astype(np.float32)).astype(np.float32) + save_case_data(name, {"input1": src, "input2": id_mat, "golden": golden}) + + elif name.startswith("mat2right"): + id_mat = np.eye(case["shape_id"][0], case["shape_id"][1], dtype=case["dtype_id"]) + src = np.random.uniform(-1.0, 1.0, size=case["shape_src"]).astype(case["dtype_src"]) + golden = np.matmul(id_mat.astype(np.float32), src.astype(np.float32)).astype(np.float32) + save_case_data(name, {"input1": id_mat, "input2": src, "golden": golden}) + + print(f"[INFO] gen_data: {name} done") \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract/launch.cpp new file mode 100644 index 0000000000..4d9cdcd789 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/launch.cpp @@ -0,0 +1,24 @@ +// 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 TEXTRACT_M2L_f16_16x16(__gm__ uint16_t *src, __gm__ uint16_t *id, __gm__ float *out); +extern "C" __global__ AICORE void TEXTRACT_M2R_f16_16x16(__gm__ uint16_t *id, __gm__ uint16_t *src, __gm__ float *out); + +void LaunchTEXTRACT_M2L_f16_16x16(uint16_t *src, uint16_t *id, float *out, void *stream) { + TEXTRACT_M2L_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)id, (__gm__ float *)out); +} + +void LaunchTEXTRACT_M2R_f16_16x16(uint16_t *id, uint16_t *src, float *out, void *stream) { + TEXTRACT_M2R_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)id, (__gm__ uint16_t *)src, (__gm__ float *)out); +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract/main.cpp new file mode 100644 index 0000000000..0c8c6e54d7 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/main.cpp @@ -0,0 +1,147 @@ +// 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 LaunchTEXTRACT_M2L_f16_16x16(uint16_t *src, uint16_t *id, float *out, void *stream); +void LaunchTEXTRACT_M2R_f16_16x16(uint16_t *id, uint16_t *src, float *out, void *stream); + +using LaunchFn = void (*)(uint16_t *, uint16_t *, float *, void *); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t input1Rows; + size_t input1Cols; + size_t input2Rows; + size_t input2Cols; + size_t outRows; + size_t outCols; +}; + +static const TestCase kCases[] = { + {"mat2left_f16_16x16", LaunchTEXTRACT_M2L_f16_16x16, 16, 16, 16, 16, 16, 16}, + {"mat2right_f16_16x16", LaunchTEXTRACT_M2R_f16_16x16, 16, 16, 16, 16, 16, 16}, +}; +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 i1Elems = tc.input1Rows * tc.input1Cols; + const size_t i2Elems = tc.input2Rows * tc.input2Cols; + const size_t outElems = tc.outRows * tc.outCols; + const size_t i1Bytes = i1Elems * sizeof(uint16_t); + const size_t i2Bytes = i2Elems * sizeof(uint16_t); + const size_t outBytes = outElems * sizeof(float); + size_t i1FileSize = i1Bytes; + size_t i2FileSize = i2Bytes; + + std::printf( + "[INFO] === case: %s (i1=%zux%zu, i2=%zux%zu, out=%zux%zu) ===\n", + tc.name, tc.input1Rows, tc.input1Cols, tc.input2Rows, tc.input2Cols, tc.outRows, tc.outCols + ); + + std::string caseDir = std::string("./") + tc.name; + + void *i1Host = nullptr; + void *i2Host = nullptr; + void *outHost = nullptr; + void *i1Device = nullptr; + void *i2Device = nullptr; + void *outDevice = nullptr; + + aclrtMallocHost(&i1Host, i1Bytes); + aclrtMallocHost(&i2Host, i2Bytes); + aclrtMallocHost(&outHost, outBytes); + + aclrtMalloc(&i1Device, i1Bytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&i2Device, i2Bytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDevice, outBytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input1.bin").c_str(), i1FileSize, i1Host, i1Bytes)) { + 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(), i2FileSize, i2Host, i2Bytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input2.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(i1Device, i1Bytes, i1Host, i1Bytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(i2Device, i2Bytes, i2Host, i2Bytes, ACL_MEMCPY_HOST_TO_DEVICE); + + tc.launch( + static_cast(i1Device), + static_cast(i2Device), + static_cast(outDevice), + stream + ); + + aclrtSynchronizeStream(stream); + aclrtMemcpy(outHost, outBytes, outDevice, outBytes, ACL_MEMCPY_DEVICE_TO_HOST); + } + + if (rc == 0 && !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 (i1Device != nullptr) aclrtFree(i1Device); + if (i2Device != nullptr) aclrtFree(i2Device); + if (outDevice != nullptr) aclrtFree(outDevice); + if (i1Host != nullptr) aclrtFreeHost(i1Host); + if (i2Host != nullptr) aclrtFreeHost(i2Host); + if (outHost != nullptr) 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/textract/textract.pto b/test/tilelang_st/npu/a5/src/st/testcase/textract/textract.pto new file mode 100644 index 0000000000..2a570b49ce --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract/textract.pto @@ -0,0 +1,176 @@ +// 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 ST kernels for textract (Mat->Left, Mat->Right) paths. +// +// Uses pto.textract for the tested path and mte surfaces for the readback path. +// Golden = src (because id x src = src for M2R, src x id = src for M2L). + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + // ---- Mat->Left: textract(src_mat -> left), mte(id_mat -> right), tmatmul -> acc -> GM ---- + func.func @TEXTRACT_M2L_f16_16x16(%src_gm: !pto.ptr, %id_gm: !pto.ptr, %out_gm: !pto.ptr) attributes {pto.aicore} { + %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 + %c0_index = arith.constant 0 : index + %false = arith.constant false + + %src_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %id_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %left_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %right_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %acc_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %src_l1 = pto.tile_buf_addr %src_tile + : !pto.tile_buf + -> !pto.ptr + %id_l1 = pto.tile_buf_addr %id_tile + : !pto.tile_buf + -> !pto.ptr + %right_ptr = pto.tile_buf_addr %right_tile + : !pto.tile_buf + -> !pto.ptr + %acc_ptr = pto.tile_buf_addr %acc_tile + : !pto.tile_buf + -> !pto.ptr + + // GM -> L1: load src (textract target) + pto.mte_gm_l1_frac %src_gm, %src_l1, 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 + + // GM -> L1: load identity (readback operand) + pto.mte_gm_l1_frac %id_gm, %id_l1, 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"] + + // textract Mat->Left: src_mat -> left (offset=0) + pto.textract ins(%src_tile, %c0_index, %c0_index : !pto.tile_buf, index, index) + outs(%left_tile : !pto.tile_buf) + + // mte L1->L0B: identity -> right (readback via proven mte surface) + pto.mte_l1_l0b %id_l1, %right_ptr, %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(%left_tile, %right_tile : !pto.tile_buf, + !pto.tile_buf) + outs(%acc_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.mte_l0c_gm %acc_ptr, %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 + } + + // ---- Mat->Right: mte(id_mat -> left), textract(src_mat -> right), tmatmul -> acc -> GM ---- + func.func @TEXTRACT_M2R_f16_16x16(%id_gm: !pto.ptr, %src_gm: !pto.ptr, %out_gm: !pto.ptr) attributes {pto.aicore} { + %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 + %c0_index = arith.constant 0 : index + %false = arith.constant false + + %id_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %src_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %left_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %right_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %acc_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %id_l1 = pto.tile_buf_addr %id_tile + : !pto.tile_buf + -> !pto.ptr + %src_l1 = pto.tile_buf_addr %src_tile + : !pto.tile_buf + -> !pto.ptr + %left_ptr = pto.tile_buf_addr %left_tile + : !pto.tile_buf + -> !pto.ptr + %acc_ptr = pto.tile_buf_addr %acc_tile + : !pto.tile_buf + -> !pto.ptr + + // GM -> L1: load identity (readback operand) + pto.mte_gm_l1_frac %id_gm, %id_l1, 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"] + + // mte L1->L0A: identity -> left (readback via proven mte surface) + pto.mte_l1_l0a %id_l1, %left_ptr, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + // GM -> L1: load src (textract target) + pto.mte_gm_l1_frac %src_gm, %src_l1, 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"] + + // textract Mat->Right: src_mat -> right (offset=0) + pto.textract ins(%src_tile, %c0_index, %c0_index : !pto.tile_buf, index, index) + outs(%right_tile : !pto.tile_buf) + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%left_tile, %right_tile : !pto.tile_buf, + !pto.tile_buf) + outs(%acc_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.mte_l0c_gm %acc_ptr, %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 + } + + } \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/CMakeLists.txt new file mode 100644 index 0000000000..59e4cdc164 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/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(textract_fp PTO_LEVEL level3) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/cases.py new file mode 100644 index 0000000000..24db5eb4e6 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/cases.py @@ -0,0 +1,25 @@ +# 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 + + +CASES = [ + { + "name": "fp_f32_f16_16x16", + "dtype_src": np.float16, + "dtype_scaling": np.float32, + "dtype_out": np.float32, + "shape_src": (16, 16), + "shape_scaling": (16, 16), + "shape_out": (16, 16), + "eps": 1e-2, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/compare.py new file mode 100644 index 0000000000..0a4b11f0d5 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/compare.py @@ -0,0 +1,45 @@ +# 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"] + shape_out = case["shape_out"] + golden = np.fromfile(os.path.join(case_dir, "golden.bin"), dtype=np.float32).reshape(shape_out) + output = np.fromfile(os.path.join(case_dir, "output.bin"), dtype=np.float32).reshape(shape_out) + + 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/textract_fp/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/gen_data.py new file mode 100644 index 0000000000..c028fa552d --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/gen_data.py @@ -0,0 +1,28 @@ +# 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) + name = case["name"] + + if name.startswith("fp"): + a = np.random.uniform(-1.0, 1.0, size=case["shape_src"]).astype(case["dtype_src"]) + b = np.random.uniform(-1.0, 1.0, size=case["shape_src"]).astype(case["dtype_src"]) + fb = np.ones(case["shape_scaling"], dtype=case["dtype_scaling"]) + golden = np.matmul(a.astype(np.float32), b.astype(np.float32)).astype(np.float32) + save_case_data(name, {"input1": a, "input2": b, "input3": fb, "golden": golden}) + + print(f"[INFO] gen_data: {name} done") \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/launch.cpp new file mode 100644 index 0000000000..ca7e708c6e --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/launch.cpp @@ -0,0 +1,19 @@ +// 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 TEXTRACT_FP_f32_f16_16x16(__gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ float *fb, __gm__ float *out); + +void LaunchTEXTRACT_FP_f32_f16_16x16(uint16_t *a, uint16_t *b, float *fb, float *out, void *stream) { + TEXTRACT_FP_f32_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ float *)fb, (__gm__ float *)out); +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/main.cpp new file mode 100644 index 0000000000..6cd3e6c382 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/main.cpp @@ -0,0 +1,162 @@ +// 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 LaunchTEXTRACT_FP_f32_f16_16x16(uint16_t *a, uint16_t *b, float *fb, float *out, void *stream); + +using LaunchFn = void (*)(uint16_t *, uint16_t *, float *, float *, void *); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t aRows; + size_t aCols; + size_t bRows; + size_t bCols; + size_t fbRows; + size_t fbCols; + size_t outRows; + size_t outCols; +}; + +static const TestCase kCases[] = { + {"fp_f32_f16_16x16", LaunchTEXTRACT_FP_f32_f16_16x16, 16, 16, 16, 16, 16, 16, 16, 16}, +}; +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.aRows * tc.aCols; + const size_t bElems = tc.bRows * tc.bCols; + const size_t fbElems = tc.fbRows * tc.fbCols; + const size_t outElems = tc.outRows * tc.outCols; + const size_t aBytes = aElems * sizeof(uint16_t); + const size_t bBytes = bElems * sizeof(uint16_t); + const size_t fbBytes = fbElems * sizeof(float); + const size_t outBytes = outElems * sizeof(float); + size_t aFileSize = aBytes; + size_t bFileSize = bBytes; + size_t fbFileSize = fbBytes; + + std::printf( + "[INFO] === case: %s (a=%zux%zu, b=%zux%zu, fb=%zux%zu, out=%zux%zu) ===\n", + tc.name, tc.aRows, tc.aCols, tc.bRows, tc.bCols, tc.fbRows, tc.fbCols, tc.outRows, tc.outCols + ); + + std::string caseDir = std::string("./") + tc.name; + + void *aHost = nullptr; + void *bHost = nullptr; + void *fbHost = nullptr; + void *outHost = nullptr; + void *aDevice = nullptr; + void *bDevice = nullptr; + void *fbDevice = nullptr; + void *outDevice = nullptr; + + aclrtMallocHost(&aHost, aBytes); + aclrtMallocHost(&bHost, bBytes); + aclrtMallocHost(&fbHost, fbBytes); + aclrtMallocHost(&outHost, outBytes); + + aclrtMalloc(&aDevice, aBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&bDevice, bBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&fbDevice, fbBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDevice, 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 && !ReadFile((caseDir + "/input3.bin").c_str(), fbFileSize, fbHost, fbBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input3.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(aDevice, aBytes, aHost, aBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(bDevice, bBytes, bHost, bBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(fbDevice, fbBytes, fbHost, fbBytes, ACL_MEMCPY_HOST_TO_DEVICE); + + tc.launch( + static_cast(aDevice), + static_cast(bDevice), + static_cast(fbDevice), + static_cast(outDevice), + stream + ); + + aclrtSynchronizeStream(stream); + aclrtMemcpy(outHost, outBytes, outDevice, outBytes, ACL_MEMCPY_DEVICE_TO_HOST); + } + + if (rc == 0 && !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 (aDevice != nullptr) aclrtFree(aDevice); + if (bDevice != nullptr) aclrtFree(bDevice); + if (fbDevice != nullptr) aclrtFree(fbDevice); + if (outDevice != nullptr) aclrtFree(outDevice); + if (aHost != nullptr) aclrtFreeHost(aHost); + if (bHost != nullptr) aclrtFreeHost(bHost); + if (fbHost != nullptr) aclrtFreeHost(fbHost); + if (outHost != nullptr) 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/textract_fp/textract_fp.pto b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto new file mode 100644 index 0000000000..cb0ca0515a --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_fp/textract_fp.pto @@ -0,0 +1,138 @@ +// 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 ST kernel for textract_fp (Acc→Mat with pre_quant f32→f16). +// +// Pipeline: mte_gm_l1_frac(A,B) → mte_l1_l0a/l0b → tmatmul → acc (L0C) +// mte_gm_l1_frac(fb) → mte_l1_fb → scaling +// textract_fp(acc, fp, 0, 0) → dst_mat (L1) +// mte_l0c_gm → GM (f32 acc for golden) + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_FP_f32_f16_16x16(%a_gm: !pto.ptr, %b_gm: !pto.ptr, %fb_gm: !pto.ptr, %out_gm: !pto.ptr) attributes {pto.aicore} { + %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 + %c1024_i64 = arith.constant 1024 : i64 + %c1536_i64 = arith.constant 1536 : i64 + %c0_index = arith.constant 0 : index + %false = arith.constant false + + // --- tile allocations --- + %a_mat = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %b_mat = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %fb_mat = pto.alloc_tile addr = %c1024_i64 + : !pto.tile_buf + + %left_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %right_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %acc_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %fp_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + // --- pointers --- + %a_l1 = pto.tile_buf_addr %a_mat + : !pto.tile_buf + -> !pto.ptr + %b_l1 = pto.tile_buf_addr %b_mat + : !pto.tile_buf + -> !pto.ptr + %fb_l1 = pto.tile_buf_addr %fb_mat + : !pto.tile_buf + -> !pto.ptr + %left_ptr = pto.tile_buf_addr %left_tile + : !pto.tile_buf + -> !pto.ptr + %right_ptr = pto.tile_buf_addr %right_tile + : !pto.tile_buf + -> !pto.ptr + %acc_ptr = pto.tile_buf_addr %acc_tile + : !pto.tile_buf + -> !pto.ptr + %fp_ptr = pto.tile_buf_addr %fp_tile + : !pto.tile_buf + -> !pto.ptr + + // --- GM -> L1: load A, B --- + pto.mte_gm_l1_frac %a_gm, %a_l1, 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.mte_gm_l1_frac %b_gm, %b_l1, 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"] + + // --- L1 -> L0A/L0B --- + pto.mte_l1_l0a %a_l1, %left_ptr, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + pto.mte_l1_l0b %b_l1, %right_ptr, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + // --- tmatmul -> acc --- + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%left_tile, %right_tile : !pto.tile_buf, + !pto.tile_buf) + outs(%acc_tile : !pto.tile_buf) + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID0"] + + // --- GM -> L1: load fb scaling params (16x16 f32 = 1 scaling per column) --- + pto.mte_gm_l1_frac %fb_gm, %fb_l1, nd2nz, + shape(%c16_i64, %c16_i64), + src_layout(%c64_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_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_FIX", "EVENT_ID1"] + + // --- L1 -> fb: load scaling to Scaling buffer --- + pto.mte_l1_fb %fb_l1, %fp_ptr, %c16_i64 + nburst(%c1_i64, %c0_i64, %c0_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + + // --- textract_fp: acc(f32) -> mat(f16) with pre_quant --- + pto.textract_fp ins(%acc_tile, %fp_tile, %c0_index, %c0_index : + !pto.tile_buf, + !pto.tile_buf, index, index) + outs(%dst_mat : !pto.tile_buf) + + // --- acc -> GM: f32 golden reference --- + pto.mte_l0c_gm %acc_ptr, %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 + } + + } \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/CMakeLists.txt new file mode 100644 index 0000000000..0e7ef24aef --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/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(textract_v2v PTO_LEVEL level3) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/cases.py new file mode 100644 index 0000000000..7355b2e12e --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/cases.py @@ -0,0 +1,23 @@ +# 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 + + +CASES = [ + { + "name": "v2v_f32_16x16", + "dtype_src": np.float32, + "dtype_out": np.float32, + "shape_src": (16, 16), + "shape_out": (16, 16), + "eps": 1e-6, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/compare.py new file mode 100644 index 0000000000..7c9b50d379 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/compare.py @@ -0,0 +1,45 @@ +# 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"] + shape_out = case["shape_out"] + golden = np.fromfile(os.path.join(case_dir, "golden.bin"), dtype=case["dtype_out"]).reshape(shape_out) + output = np.fromfile(os.path.join(case_dir, "output.bin"), dtype=case["dtype_out"]).reshape(shape_out) + + 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/textract_v2v/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/gen_data.py new file mode 100644 index 0000000000..231dbba10a --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/gen_data.py @@ -0,0 +1,26 @@ +# 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) + name = case["name"] + + if name.startswith("v2v"): + src = np.random.uniform(-1.0, 1.0, size=case["shape_src"]).astype(case["dtype_src"]) + golden = src.copy() + save_case_data(name, {"input1": src, "golden": golden}) + + print(f"[INFO] gen_data: {name} done") \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/launch.cpp new file mode 100644 index 0000000000..e5fbbbd35b --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/launch.cpp @@ -0,0 +1,19 @@ +// 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 TEXTRACT_V2V_ND_f32_16x16(__gm__ float *src, __gm__ float *out); + +void LaunchTEXTRACT_V2V_ND_f32_16x16(float *src, float *out, void *stream) { + TEXTRACT_V2V_ND_f32_16x16<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)out); +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/main.cpp new file mode 100644 index 0000000000..6e11102966 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/main.cpp @@ -0,0 +1,128 @@ +// 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 LaunchTEXTRACT_V2V_ND_f32_16x16(float *src, float *out, void *stream); + +using LaunchFn = void (*)(float *, float *, void *); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t srcRows; + size_t srcCols; + size_t outRows; + size_t outCols; +}; + +static const TestCase kCases[] = { + {"v2v_f32_16x16", LaunchTEXTRACT_V2V_ND_f32_16x16, 16, 16, 16, 16}, +}; +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.srcRows * tc.srcCols; + const size_t outElems = tc.outRows * tc.outCols; + const size_t srcBytes = srcElems * sizeof(float); + const size_t outBytes = outElems * sizeof(float); + size_t srcFileSize = srcBytes; + + std::printf( + "[INFO] === case: %s (src=%zux%zu, out=%zux%zu) ===\n", + tc.name, tc.srcRows, tc.srcCols, tc.outRows, tc.outCols + ); + + std::string caseDir = std::string("./") + tc.name; + + void *srcHost = nullptr; + void *outHost = nullptr; + void *srcDevice = nullptr; + void *outDevice = nullptr; + + aclrtMallocHost(&srcHost, srcBytes); + aclrtMallocHost(&outHost, outBytes); + + aclrtMalloc(&srcDevice, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDevice, outBytes, 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) { + aclrtMemcpy(srcDevice, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE); + + tc.launch( + static_cast(srcDevice), + static_cast(outDevice), + stream + ); + + aclrtSynchronizeStream(stream); + aclrtMemcpy(outHost, outBytes, outDevice, outBytes, ACL_MEMCPY_DEVICE_TO_HOST); + } + + if (rc == 0 && !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 (srcDevice != nullptr) aclrtFree(srcDevice); + if (outDevice != nullptr) aclrtFree(outDevice); + if (srcHost != nullptr) aclrtFreeHost(srcHost); + if (outHost != nullptr) 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/textract_v2v/textract_v2v.pto b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/textract_v2v.pto new file mode 100644 index 0000000000..4c950713c6 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/textract_v2v/textract_v2v.pto @@ -0,0 +1,59 @@ +// 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 ST kernel for textract Vec->Vec ND (UB sub-window copy). +// +// kernel_kind=vector (pure vec path, no cube needed). +// Pipeline: tload(src) → textract(src_vec→dst_vec) → tstore(dst) + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + func.func @TEXTRACT_V2V_ND_f32_16x16(%src_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c256 = arith.constant 256 : index + %c0_i64 = arith.constant 0 : i64 + %c1024_i64 = arith.constant 1024 : i64 + + %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, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf32> + + %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, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf32> -> !pto.partition_tensor_view<1x1x1x16x16xf32> + + %src_vec = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec = pto.alloc_tile addr = %c1024_i64 + : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf32>) + outs(%src_vec : !pto.tile_buf) + + pto.textract ins(%src_vec, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_vec : !pto.tile_buf) + + pto.tstore ins(%dst_vec : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x16xf32>) + + return + } + +} \ No newline at end of file diff --git a/tilelang-dsl/python/tilelang_dsl/semantic.py b/tilelang-dsl/python/tilelang_dsl/semantic.py index 4ba6888e54..59862dc14a 100644 --- a/tilelang-dsl/python/tilelang_dsl/semantic.py +++ b/tilelang-dsl/python/tilelang_dsl/semantic.py @@ -7207,8 +7207,8 @@ def _require_i64_like_expr(self, expr: SemanticExpr, context: str) -> None: if isinstance(expr.type, SemanticIndexType): return scalar = self._require_scalar_expr(expr, context) - if scalar.dtype != i64: - raise TypeError(f"{context} must be an i64 or index value in TileLang DSL") + if scalar.dtype not in (i64, i32): + raise TypeError(f"{context} must be an i64, i32, or index value in TileLang DSL") def _require_tail_remaining_expr(self, expr: SemanticExpr, context: str) -> None: if isinstance(expr.type, SemanticIndexType):