diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 6903ff9c5..13cd90609 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -7475,7 +7475,7 @@ pto.tconcatidx ins(%src0, %src1, %idx0, %idx1 : **Summary:** Gathers elements from a source tile using one of three PTO-ISA-compatible forms: -- index gather: `src + indices + tmp -> dst` +- index gather: `src + indices[+ tmp] -> dst` - compare gather: `src + kValue + tmp -> dst + cdst` - mask-pattern gather: `src + maskPattern -> dst` @@ -7501,7 +7501,7 @@ Mask form: | `dst` | `pto.tile_buf` | Main destination tile | | `cdst` | `Optional` | Secondary destination tile used only by compare form | | `indices` | `Optional` | Index tile used only by index form | -| `tmp` | `Optional` | Temporary tile used by index form and compare form | +| `tmp` | `Optional` | Temporary tile used by compare form; optionally used by index form on A2/A3 (required), A5 (optional) | | `kValue` | `Optional` | Scalar compare value used only by compare form | | `maskPattern` | `Optional` | Mask pattern used only by mask form | | `cmpMode` | `Optional` | Compare mode used only by compare form; defaults to `eq` when omitted | @@ -7535,18 +7535,18 @@ pto.tgather ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf<. **Constraints & Verification:** - Exactly one of the following forms must be used: - - index form: `indices` and `tmp` + - index form: `indices` (A5: `tmp` is optional; A2/A3: `tmp` is required) - compare form: `kValue`, `tmp`, and `cdst` - mask form: `maskPattern` - **Index gather: implementation checks (A2/A3)**: - `src` and `dst` element types must match and be one of `i16/i32/f16/f32`. - `indices` element type must be `i32`. - - `tmp` element type must match `indices`. + - `tmp` is required; `tmp` element type must match `indices`. - `indices` and `tmp` must have the same valid shape. - **Index gather: implementation checks (A5)**: - `src` and `dst` element types must match and be one of `i8/i16/i32/f16/f32`, or a target-supported fp8 type (`f8E4M3*`/`f8E5M2*`). - `indices` element type must be `i16` or `i32`. - - PTO IR does not impose an extra tmp shape or valid-shape relation in the A5 index form. + - `tmp` is optional; PTO IR does not impose an extra tmp shape or valid-shape relation in the A5 index form. - **Compare gather: implementation checks (A2/A3)**: - `dst` element type must be `i32`. - `src` element type must be `f16/f32`, or `i32` when `cmpMode=eq`. diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index af6dd2859..5395cc61f 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7758,12 +7758,15 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); Type idxTy = getIndices().getType(); - Type tmpTy = getTmp().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src", allowA5ElemTypes)) || failed(verifyTileBufCommon(*this, dstTy, "dst", allowA5ElemTypes)) || - failed(verifyTileBufCommon(*this, idxTy, "indices")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + failed(verifyTileBufCommon(*this, idxTy, "indices"))) return failure(); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); + } Type srcElem = getElemTy(srcTy); Type dstElem = getElemTy(dstTy); @@ -7807,10 +7810,10 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { } if (!allowA5ElemTypes) { - Type tmpElem = getElemTy(tmpTy); + Type tmpElem = getElemTy(getTmp().getType()); if (tmpElem != idxElem) return emitOpError("expects tmp and indices to have the same element type"); - if (failed(verifyTileBufSameValidShape(*this, idxTy, tmpTy, "indices", "tmp"))) + if (failed(verifyTileBufSameValidShape(*this, idxTy, getTmp().getType(), "indices", "tmp"))) return failure(); } return success(); @@ -7898,8 +7901,8 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { return emitOpError("compare-form tgather does not take indices"); return verifyCompareForm(/*allowA5SrcTypes=*/true); } - if (!getIndices() || !getTmp()) - return emitOpError("index-form tgather expects both indices and tmp"); + if (!getIndices()) + return emitOpError("index-form tgather expects indices"); return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true); }; diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 33e5e6643..4abbf0622 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -10111,15 +10111,17 @@ struct PTOGatherToEmitC : public OpConversionPattern { return rewriter.notifyMatchFailure(op, (name + " must be emitc::OpaqueType (tile)").str()); }; - // Case 1: index-based TGATHER(dst, src0, indices, tmp) + // Case 1: index-based TGATHER(dst, src0, indices[, tmp]) if (Value idx = adaptor.getIndices()) { idx = peelUnrealized(idx); - Value tmp = peelUnrealized(adaptor.getTmp()); + SmallVector operands{dst, src0, idx}; + if (Value tmp = adaptor.getTmp()) + operands.push_back(peelUnrealized(tmp)); rewriter.create( loc, TypeRange{}, "TGATHER", /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/ValueRange{dst, src0, idx, tmp}); + /*operands=*/operands); rewriter.eraseOp(op); return success(); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 46da6fb59..efbaa41eb 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -3654,7 +3654,8 @@ struct PTOViewToMemrefPass } if (maskPattern) { - rewriter.replaceOpWithNewOp( + replaceOpWithClonedAttrs( + rewriter, op, TypeRange{}, src, @@ -3678,7 +3679,8 @@ struct PTOViewToMemrefPass return; } - rewriter.replaceOpWithNewOp( + replaceOpWithClonedAttrs( + rewriter, op, TypeRange{}, src, @@ -3693,23 +3695,33 @@ struct PTOViewToMemrefPass continue; } - if (indices || tmp) { + if (indices) { auto indicesTy = dyn_cast(indices.getType()); - auto tmpTy = dyn_cast(tmp.getType()); - if (!indicesTy || !tmpTy) { - op.emitError("index-form tgather expects indices/tmp to be memref yet"); + if (!indicesTy) { + op.emitError("index-form tgather expects indices to be memref yet"); signalPassFailure(); return; } + Value tmpVal; + if (tmp) { + auto tmpTy = dyn_cast(tmp.getType()); + if (!tmpTy) { + op.emitError("index-form tgather expects tmp to be memref yet"); + signalPassFailure(); + return; + } + tmpVal = tmp; + } - rewriter.replaceOpWithNewOp( + replaceOpWithClonedAttrs( + rewriter, op, TypeRange{}, src, dst, /*cdst=*/Value(), indices, - tmp, + tmpVal, /*kValue=*/Value(), /*maskPattern=*/pto::MaskPatternAttr(), /*cmpMode=*/pto::CmpModeAttr(), diff --git a/ptodsl/ptodsl/tilelib/templates/__init__.py b/ptodsl/ptodsl/tilelib/templates/__init__.py index eee2fa0ae..73bcfb66c 100644 --- a/ptodsl/ptodsl/tilelib/templates/__init__.py +++ b/ptodsl/ptodsl/tilelib/templates/__init__.py @@ -48,6 +48,7 @@ ("a5", "pto.tfillpad"): ".a5.tfillpad", ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand", ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace", + ("a5", "pto.tgather"): ".a5.tgather", ("a5", "pto.tgatherb"): ".a5.tgatherb", ("a5", "pto.tgemv"): ".a5.tgemv", ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc", diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgather.py b/ptodsl/ptodsl/tilelib/templates/a5/tgather.py new file mode 100644 index 000000000..86f2ba472 --- /dev/null +++ b/ptodsl/ptodsl/tilelib/templates/a5/tgather.py @@ -0,0 +1,212 @@ +# 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. + +"""PTODSL TileLib templates for pto.tgather.""" + +from ptodsl import pto +import ptodsl.tilelib as tilelib +from ._common import NUMERIC_DTYPES +from ._common import element_store_dist + + +def gather_dtype_signatures(dtypes=NUMERIC_DTYPES): + res = [] + for dtype in dtypes: + for dtype_indices in ('i16', 'ui16', "i32", "ui32"): + res.append((dtype, dtype, dtype_indices)) + return res + + +_GATHER_MASK_DTYPES = [(d, d) for d in NUMERIC_DTYPES] + +_2X_MASK_PATTERNS = ("P0101", "P1010") +_4X_MASK_PATTERNS = ("P0001", "P0010", "P0100", "P1000") + + +def _is_2x_mask(mask_pattern=None, **_): + return mask_pattern in _2X_MASK_PATTERNS + + +def _is_4x_mask(mask_pattern=None, **_): + return mask_pattern in _4X_MASK_PATTERNS + + +def _is_p1111(mask_pattern=None, **_): + return mask_pattern == "P1111" + + +@tilelib.tile_template( + op="pto.tgather", + target="a5", + name="template_tgather", + dtypes=gather_dtype_signatures(), + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False +) +def template_tgather( + src: pto.Tile, + dst: pto.Tile, + indices: pto.Tile): + dtype = dst.element_type + dtype_indices = indices.element_type + elem_bytes = pto.bytewidth(dtype) + elem_bytes_s1 = pto.bytewidth(dtype_indices) + valid_rows, valid_cols = indices.valid_shape + src_ptr = src.as_ptr() + if pto.const_expr(elem_bytes == 2 and elem_bytes_s1 == 4): + indices_type = pto.ui32 + lanes = pto.elements_per_vreg(indices_type) + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + indices_reg = pto.vlds(indices[row, col:]) + mask, remained = pto.make_mask(indices_type, remained) + dst_reg = pto.vgather2_bc(src_ptr, pto.vbitcast(indices_reg, indices_type), mask) + pto.vsts(dst_reg, dst[row, col:], mask) + elif pto.const_expr(elem_bytes == 1): + packed_lanes = pto.elements_per_vreg(dtype) >> 1 + result_elem = pto.ui16 if pto.const_expr(dtype == pto.ui8) else pto.i16 + result_ty = pto.vreg_type(packed_lanes, result_elem) + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, packed_lanes): + indices_reg = pto.vlds(indices[row, col:]) + mask, remained = pto.make_mask(result_elem, remained) + dst_reg = pto.vgather2(src_ptr, pto.vbitcast(indices_reg, pto.ui16), mask, result_vreg_type=result_ty) + pto.vsts(dst_reg, dst[row, col:], mask, dist=pto.VStoreDist.PK_B16) + elif pto.const_expr(elem_bytes == 4): + indices_type = pto.ui32 + lanes = pto.elements_per_vreg(indices_type) + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + indices_reg = pto.vlds(indices[row, col:]) + mask, remained = pto.make_mask(indices_type, remained) + dst_reg = pto.vgather2(src_ptr, pto.vbitcast(indices_reg, indices_type), mask) + pto.vsts(dst_reg, dst[row, col:], mask) + else: + indices_type = pto.ui16 + lanes = pto.elements_per_vreg(indices_type) + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + indices_reg = pto.vlds(indices[row, col:]) + mask, remained = pto.make_mask(indices_type, remained) + dst_reg = pto.vgather2(src_ptr, pto.vbitcast(indices_reg, indices_type), mask) + pto.vsts(dst_reg, dst[row, col:], mask) + + +@tilelib.tile_template( + op="pto.tgather", + target="a5", + name="template_tgather_mask_2x", + dtypes=_GATHER_MASK_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + tags=("gather", "mask"), + constraints=(_is_2x_mask,), + id=1, +) +def template_tgather_mask_2x(src: pto.Tile, dst: pto.Tile): + dtype = dst.element_type + lanes = pto.elements_per_vreg(dtype) + valid_rows = src.valid_shape[0] + dst_valid_cols = dst.valid_shape[1] + mask_pattern = pto.get_op_attr("mask_pattern") + for row in range(0, valid_rows, 1): + remained = dst_valid_cols + for dst_col in range(0, dst_valid_cols, lanes): + src_col = dst_col * 2 + src_reg0 = pto.vlds(src[row, src_col:]) + src_reg1 = pto.vlds(src[row, src_col + lanes:]) + even, odd = pto.vdintlv(src_reg0, src_reg1) + if mask_pattern == "P0101": + result = even + else: + result = odd + mask, remained = pto.make_mask(dtype, remained) + pto.vsts(result, dst[row, dst_col:], mask) + + +@tilelib.tile_template( + op="pto.tgather", + target="a5", + name="template_tgather_mask_4x", + dtypes=_GATHER_MASK_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + tags=("gather", "mask"), + constraints=(_is_4x_mask,), + id=2, +) +def template_tgather_mask_4x(src: pto.Tile, dst: pto.Tile): + dtype = dst.element_type + lanes = pto.elements_per_vreg(dtype) + valid_rows = src.valid_shape[0] + dst_valid_cols = dst.valid_shape[1] + mask_pattern = pto.get_op_attr("mask_pattern") + for row in range(0, valid_rows, 1): + remained = dst_valid_cols + for dst_col in range(0, dst_valid_cols, lanes): + src_col = dst_col * 4 + src_reg0 = pto.vlds(src[row, src_col:]) + src_reg1 = pto.vlds(src[row, src_col + lanes:]) + src_reg2 = pto.vlds(src[row, src_col + 2 * lanes:]) + src_reg3 = pto.vlds(src[row, src_col + 3 * lanes:]) + e0, o0 = pto.vdintlv(src_reg0, src_reg1) + e1, o1 = pto.vdintlv(src_reg2, src_reg3) + if mask_pattern == "P0001": + result, _ = pto.vdintlv(e0, e1) + elif mask_pattern == "P0010": + result, _ = pto.vdintlv(o0, o1) + elif mask_pattern == "P0100": + _, result = pto.vdintlv(e0, e1) + else: + _, result = pto.vdintlv(o0, o1) + mask, remained = pto.make_mask(dtype, remained) + pto.vsts(result, dst[row, dst_col:], mask) + + +@tilelib.tile_template( + op="pto.tgather", + target="a5", + name="template_tgather_mask_p1111", + dtypes=_GATHER_MASK_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + tags=("gather", "mask"), + constraints=(_is_p1111,), + id=3, +) +def template_tgather_mask_p1111(src: pto.Tile, dst: pto.Tile): + dtype = dst.element_type + lanes = pto.elements_per_vreg(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): + src_reg = pto.vlds(src[row, col:]) + mask, remained = pto.make_mask(dtype, remained) + pto.vsts(src_reg, dst[row, col:], mask) + diff --git a/test/lit/pto/tgather_a2a3_index_emitc.pto b/test/lit/pto/tgather_a2a3_index_emitc.pto new file mode 100644 index 000000000..2dd05e0c0 --- /dev/null +++ b/test/lit/pto/tgather_a2a3_index_emitc.pto @@ -0,0 +1,15 @@ +// RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s + +module { + func.func @tgather_a2a3_index_with_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: TGATHER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); diff --git a/test/lit/pto/tgather_a2a3_index_invalid.pto b/test/lit/pto/tgather_a2a3_index_invalid.pto new file mode 100644 index 000000000..9cf08b0a1 --- /dev/null +++ b/test/lit/pto/tgather_a2a3_index_invalid.pto @@ -0,0 +1,45 @@ +// RUN: not ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s + +module { + func.func @tgather_a2a3_index_missing_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a2a3_index_missing_indices() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a2a3_index_invalid_src_dtype() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a2a3_index_invalid_indices_dtype() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: error: 'pto.tgather' op index-form tgather expects both indices and tmp +// CHECK: error: 'pto.tgather' op index-form tgather expects both indices and tmp +// CHECK: error: 'pto.tgather' op expects gather src/dst element type to be i16/i32/f16/f32 +// CHECK: error: 'pto.tgather' op expects indices element type to be i32 diff --git a/test/lit/pto/tgather_a5_index_no_tmp_emitc.pto b/test/lit/pto/tgather_a5_index_no_tmp_emitc.pto new file mode 100644 index 000000000..e3d15cc61 --- /dev/null +++ b/test/lit/pto/tgather_a5_index_no_tmp_emitc.pto @@ -0,0 +1,25 @@ +// RUN: ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s + +module { + func.func @tgather_a5_index_no_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a5_index_with_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: TGATHER +// CHECK: TGATHER diff --git a/test/lit/pto/tgather_a5_index_optional_tmp_invalid.pto b/test/lit/pto/tgather_a5_index_optional_tmp_invalid.pto new file mode 100644 index 000000000..472a8c465 --- /dev/null +++ b/test/lit/pto/tgather_a5_index_optional_tmp_invalid.pto @@ -0,0 +1,33 @@ +// RUN: not ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s + +module { + func.func @tgather_a5_index_missing_indices() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a5_index_invalid_src_dtype() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tgather_a5_index_invalid_indices_dtype() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: error: 'pto.tgather' op index-form tgather expects indices +// CHECK: error: 'pto.tgather' op expects A5 gather src/dst element type to be i8/i16/i32/f16/f32 +// CHECK: error: 'pto.tgather' op expects indices element type to be i32 or i16 diff --git a/test/tilelib-st/a5/tgather/case.py b/test/tilelib-st/a5/tgather/case.py new file mode 100644 index 000000000..b67dd4711 --- /dev/null +++ b/test/tilelib-st/a5/tgather/case.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +# 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. + +# PTODSL rewrite of test/tilelang_st/npu/a5/src/st/testcase/tgather. +# +# Two modes are supported: +# - Indexed gather: dst elements are gathered from src at positions specified +# by the corresponding index values. +# - Masked gather: dst elements are extracted from src at positions determined +# by a mask pattern (e.g., P0101 extracts even-positioned elements at 2x +# compression). This is the inverse of masked scatter. + + +from pathlib import Path +import sys +import zlib + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +PTO_TO_NP_DTYPE = { + pto.f32: np.float32, + pto.f16: np.float16, + pto.i32: np.int32, + pto.i16: np.int16, + pto.i8: np.int8, + pto.ui32: np.uint32, + pto.ui16: np.uint16, + pto.ui8: np.uint8, +} + + +def npy_dtype(pto_type) -> np.dtype: + return PTO_TO_NP_DTYPE[pto_type] + + +# Each case is (name, src_dtype, idx_dtype, src_shape, dst_shape). +# The indices tile has the same shape as dst and contains element +# indices into the flat src tile. +CASE_SHAPES = [ + ("f32_1x128_1x64", pto.f32, pto.i32, (1, 128), (1, 64)), + ("f32_1x64_1x32", pto.f32, pto.i32, (1, 64), (1, 32)), + ("f32_i32_32x1024_16x64", pto.f32, pto.i32, (32, 1024), (16, 64)), + ("i32_i32_32x512_16x256", pto.i32, pto.i32, (32, 512), (16, 256)), + ("f16_i16_16x1024_16x128", pto.f16, pto.i16, (16, 1024), (16, 128)), + ("i16_i16_32x256_32x64", pto.i16, pto.i16, (32, 256), (32, 64)), + ("i8_i16_16x128_16x64", pto.i8, pto.i16, (16, 128), (16, 64)), + ("i8_ui16_16x128_16x64", pto.i8, pto.ui16, (16, 128), (16, 64)), + ("ui8_ui16_16x128_16x64", pto.ui8, pto.ui16, (16, 128), (16, 64)), +] + + +def _tgather_body( + src_ptr, + offset_ptr, + dst_ptr, + *, + src_rows, + src_cols, + dst_rows, + dst_cols, + src_dtype, + idx_dtype, +): + """Shared kernel body for the tgather cases. + + Loads *src* and *offset* tiles from GM, performs gather using + ``pto.tile.gather``, and stores *dst* back to GM. + """ + + src_view = pto.make_tensor_view( + src_ptr, shape=[src_rows, src_cols], strides=[src_cols, 1] + ) + indices_view = pto.make_tensor_view( + offset_ptr, shape=[dst_rows, dst_cols], strides=[dst_cols, 1] + ) + dst_view = pto.make_tensor_view( + dst_ptr, shape=[dst_rows, dst_cols], strides=[dst_cols, 1] + ) + + src_tile = pto.alloc_tile(shape=[src_rows, src_cols], dtype=src_dtype) + indices_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=idx_dtype) + dst_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=src_dtype) + + pto.tile.load(src_view, src_tile) + pto.tile.load(indices_view, indices_tile) + pto.tile.gather(src_tile, dst_tile, indices=indices_tile) + pto.tile.store(dst_tile, dst_view) + + +def _tgather_mask_body( + src_ptr, + dst_ptr, + *, + src_rows, + src_cols, + dst_rows, + dst_cols, + dtype, + pattern, +): + """Shared kernel body for masked tgather cases. + + Loads *src* tile from GM, performs masked gather using + ``pto.tile.gather`` with ``mask_pattern``, and stores *dst* back to GM. + """ + + src_view = pto.make_tensor_view( + src_ptr, shape=[src_rows, src_cols], strides=[src_cols, 1] + ) + dst_view = pto.make_tensor_view( + dst_ptr, shape=[dst_rows, dst_cols], strides=[dst_cols, 1] + ) + + src_tile = pto.alloc_tile(shape=[src_rows, src_cols], dtype=dtype) + dst_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=dtype) + + pto.tile.load(src_view, src_tile) + pto.tile.gather(src_tile, dst_tile, mask_pattern=pattern) + pto.tile.store(dst_tile, dst_view) + + +# One decorated kernel per case, each binding static shapes at definition time. +_tgather_kernels = {} +for _name, _src_dtype, _idx_dtype, _src_shape, _dst_shape in CASE_SHAPES: + _sr, _sc = _src_shape + _dr, _dc = _dst_shape + + def _make( + sr=_sr, + sc=_sc, + dr=_dr, + dc=_dc, + sdt=_src_dtype, + idt=_idx_dtype, + kernel_name=f"tgather_{_name}", + ): + @pto.jit(name=kernel_name, target="a5") + def _kernel( + src_ptr: pto.ptr(sdt, "gm"), + offset_ptr: pto.ptr(idt, "gm"), + dst_ptr: pto.ptr(sdt, "gm"), + ): + _tgather_body( + src_ptr, + offset_ptr, + dst_ptr, + src_rows=sr, + src_cols=sc, + dst_rows=dr, + dst_cols=dc, + src_dtype=sdt, + idx_dtype=idt, + ) + + return _kernel + + _tgather_kernels[_name] = _make() + + +# --- Masked gather cases --- +# TGATHER mask is the inverse of TSCATTER mask: +# - TSCATTER row: src(R, C_src) -> dst(R, C_dst), C_dst = C_src * times +# - TGATHER row: src(R, C_src) -> dst(R, C_dst), C_src = C_dst * times +# (name, dtype, src_shape, dst_shape, pattern) +MASK_CASES = [ + ("mask_f16_16x64_16x64_P1111", pto.f16, (16, 64), (16, 64), "P1111"), + ("mask_f32_16x64_16x64_P1111", pto.f32, (16, 64), (16, 64), "P1111"), + ("mask_i32_16x64_16x64_P1111", pto.i32, (16, 64), (16, 64), "P1111"), + ("mask_i16_16x64_16x64_P1111", pto.i16, (16, 64), (16, 64), "P1111"), + ("mask_f16_16x128_16x64_P1010", pto.f16, (16, 128), (16, 64), "P1010"), + ("mask_f16_16x128_16x64_P0101", pto.f16, (16, 128), (16, 64), "P0101"), + ("mask_f32_16x128_16x64_P1010", pto.f32, (16, 128), (16, 64), "P1010"), + ("mask_f32_16x128_16x64_P0101", pto.f32, (16, 128), (16, 64), "P0101"), + ("mask_i32_16x128_16x64_P1010", pto.i32, (16, 128), (16, 64), "P1010"), + ("mask_i32_16x128_16x64_P0101", pto.i32, (16, 128), (16, 64), "P0101"), + ("mask_i16_16x128_16x64_P1010", pto.i16, (16, 128), (16, 64), "P1010"), + ("mask_i16_16x128_16x64_P0101", pto.i16, (16, 128), (16, 64), "P0101"), + ("mask_f16_16x256_16x64_P1000", pto.f16, (16, 256), (16, 64), "P1000"), + ("mask_f16_16x256_16x64_P0100", pto.f16, (16, 256), (16, 64), "P0100"), + ("mask_f16_16x256_16x64_P0010", pto.f16, (16, 256), (16, 64), "P0010"), + ("mask_f16_16x256_16x64_P0001", pto.f16, (16, 256), (16, 64), "P0001"), + ("mask_f32_16x256_16x64_P1000", pto.f32, (16, 256), (16, 64), "P1000"), + ("mask_f32_16x256_16x64_P0100", pto.f32, (16, 256), (16, 64), "P0100"), + ("mask_f32_16x256_16x64_P0010", pto.f32, (16, 256), (16, 64), "P0010"), + ("mask_f32_16x256_16x64_P0001", pto.f32, (16, 256), (16, 64), "P0001"), + ("mask_i32_16x256_16x64_P1000", pto.i32, (16, 256), (16, 64), "P1000"), + ("mask_i32_16x256_16x64_P0100", pto.i32, (16, 256), (16, 64), "P0100"), + ("mask_i32_16x256_16x64_P0010", pto.i32, (16, 256), (16, 64), "P0010"), + ("mask_i32_16x256_16x64_P0001", pto.i32, (16, 256), (16, 64), "P0001"), + ("mask_i16_16x256_16x64_P1000", pto.i16, (16, 256), (16, 64), "P1000"), + ("mask_i16_16x256_16x64_P0100", pto.i16, (16, 256), (16, 64), "P0100"), + ("mask_i16_16x256_16x64_P0010", pto.i16, (16, 256), (16, 64), "P0010"), + ("mask_i16_16x256_16x64_P0001", pto.i16, (16, 256), (16, 64), "P0001"), +] + +_mask_kernels = {} +for _name, _dtype, _src_shape, _dst_shape, _pattern in MASK_CASES: + _sr, _sc = _src_shape + _dr, _dc = _dst_shape + + def _make_mask( + sr=_sr, + sc=_sc, + dr=_dr, + dc=_dc, + dtype=_dtype, + pattern=_pattern, + kernel_name=f"tgather_{_name}", + ): + @pto.jit(name=kernel_name, target="a5") + def _kernel( + src_ptr: pto.ptr(dtype, "gm"), + dst_ptr: pto.ptr(dtype, "gm"), + ): + _tgather_mask_body( + src_ptr, + dst_ptr, + src_rows=sr, + src_cols=sc, + dst_rows=dr, + dst_cols=dc, + dtype=dtype, + pattern=pattern, + ) + + return _kernel + + _mask_kernels[_name] = _make_mask() + + +def _make_inputs(name, src_dtype, idx_dtype, src_shape, dst_shape): + src_np = npy_dtype(src_dtype) + idx_np = npy_dtype(idx_dtype) + rng_seed = zlib.crc32(name.encode("utf-8")) & 0xFFFFFFFF + rng = np.random.RandomState(rng_seed) + + if np.issubdtype(src_np, np.floating): + src = rng.uniform(-10.0, 10.0, size=src_shape).astype(src_np) + else: + src = rng.randint(-20, 20, size=src_shape).astype(src_np) + + num_src_elements = int(np.prod(src_shape)) + raw_offsets = rng.randint(0, num_src_elements, size=dst_shape).astype(idx_np) + + return [src, raw_offsets] + + +def _make_expected(src, offsets): + flat_src = src.reshape(-1) + dst = np.empty(offsets.shape, dtype=src.dtype) + flat_offsets = offsets.ravel() + flat_dst = dst.ravel() + for i in range(len(flat_dst)): + flat_dst[i] = flat_src[int(flat_offsets[i])] + return dst + + +_PATTERN_PARAMS = { + "P0101": (0, 2), + "P1010": (1, 2), + "P0001": (0, 4), + "P0010": (1, 4), + "P0100": (2, 4), + "P1000": (3, 4), + "P1111": (0, 1), +} + + +def _make_mask_inputs(name, dtype, src_shape): + np_dt = npy_dtype(dtype) + rng = np.random.RandomState(zlib.crc32(name.encode("utf-8")) & 0xFFFFFFFF) + if np.issubdtype(np_dt, np.floating): + src = rng.uniform(-10.0, 10.0, size=src_shape).astype(np_dt) + else: + src = rng.randint(-20, 20, size=src_shape).astype(np_dt) + return [src] + + +def _gather_mask_golden(src, dst_rows, dst_cols, pattern): + offset, stride = _PATTERN_PARAMS[pattern] + dst = np.empty((dst_rows, dst_cols), dtype=src.dtype) + for i in range(dst_rows): + for j in range(dst_cols): + dst[i, j] = src[i, j * stride + offset] + return dst + + +CASES = [] +for _name, _src_dtype, _idx_dtype, _src_shape, _dst_shape in CASE_SHAPES: + CASES.append( + golden_output_case( + "tgather_" + _name, + _tgather_kernels[_name], + inputs=lambda _n=_name, _sd=_src_dtype, _id=_idx_dtype, _ss=_src_shape, _ds=_dst_shape: ( + _make_inputs(_n, _sd, _id, _ss, _ds) + ), + expected=_make_expected, + rtol=1e-6, + atol=1e-6, + ) + ) + +for _name, _dtype, _src_shape, _dst_shape, _pattern in MASK_CASES: + CASES.append( + golden_output_case( + "tgather_" + _name, + _mask_kernels[_name], + inputs=lambda _n=_name, _d=_dtype, _ss=_src_shape: ( + _make_mask_inputs(_n, _d, _ss) + ), + expected=lambda src, _dr=_dst_shape[0], _dc=_dst_shape[1], _p=_pattern: ( + _gather_mask_golden(src, _dr, _dc, _p) + ), + rtol=1e-6, + atol=1e-6, + ) + ) + + +auto_main(globals())