From ad3e172b988e23b38974ce95ea526c8b8ff7c16f Mon Sep 17 00:00:00 2001 From: frank-deng <13382084679@163.com> Date: Mon, 20 Jul 2026 09:56:23 +0800 Subject: [PATCH 1/2] add TGATHERB --- lib/PTO/IR/VPTO.cpp | 9 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 8 +- .../docs/user_guide/08-compute-operations.md | 34 +++- ptodsl/ptodsl/_ops.py | 9 + ptodsl/ptodsl/_tile_namespace.py | 1 + ptodsl/ptodsl/tilelib/templates/__init__.py | 1 + .../ptodsl/tilelib/templates/a5/tgatherb.py | 50 +++++ test/tilelib-st/a5/tgatherb/case.py | 171 ++++++++++++++++++ 8 files changed, 273 insertions(+), 10 deletions(-) create mode 100644 ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py create mode 100644 test/tilelib-st/a5/tgatherb/case.py diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 20769c2de0..c62247f21b 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -4701,8 +4701,13 @@ LogicalResult VgatherbOp::verify() { return emitOpError("offset vector must use integer element type"); if (offsetsElemType.getWidth() != 32) return emitOpError("currently requires 32-bit offset vector elements"); - if (offsetsType.getElementCount() != resultType.getElementCount()) - return emitOpError("offset and result vectors must have the same element count"); + // vgatherb is a 32-byte block gather: each offset addresses one 32-byte block. + // The offset vector holds VL/32 block addresses (always ui32), while the + // result vector holds VL/sizeof(T) elements of the data type. These counts + // only coincide when sizeof(T)==4 (e.g. f32/i32/ui32). For smaller types + // the result has more elements than the offset, which is correct because the + // hardware interprets the low VL/32 bytes of the offset register as block + // addresses and gathers VL bytes of data per invocation. return success(); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 43593f4cb7..5ec948c025 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -3739,12 +3739,8 @@ struct PTOViewToMemrefPass return; } - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - offsets, - dst); + replaceOpWithClonedAttrs( + rewriter, op, TypeRange{}, src, offsets, dst); } DefaultInlineVector logops; diff --git a/ptodsl/docs/user_guide/08-compute-operations.md b/ptodsl/docs/user_guide/08-compute-operations.md index 5182f69e44..65965c71b2 100644 --- a/ptodsl/docs/user_guide/08-compute-operations.md +++ b/ptodsl/docs/user_guide/08-compute-operations.md @@ -248,8 +248,38 @@ pto.tile.gather(tmp_sort_tile, top_scores, mask_pattern="P0101") pto.tile.gather(tmp_sort_tile, top_indices, mask_pattern="P1010") ``` -The low-level aliases `pto.tsort32`, `pto.tmrgsort`, and `pto.tgather` are also -available when a kernel needs to bypass the `pto.tile` namespace. +#### `pto.tile.gatherb(src: Tile, offsets: Tile, dst: Tile) -> None` + +**Description**: Gathers elements from a source tile into a destination tile +using byte offsets. Each element of `offsets` is a byte address into the flat +byte representation of `src`; the element (or block of elements starting at that +byte position) is written into the corresponding position of `dst`. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `src` | `Tile` | Source tile containing the data to gather from | +| `offsets` | `Tile` | Byte offset tile (`ui32` dtype) specifying byte positions in `src` | +| `dst` | `Tile` | Destination tile (must use row-major layout; element size 1, 2, or 4 bytes) | + +**Returns**: None + +**Constraints**: + +- `offsets` must have `ui32` dtype. +- `dst` must use row-major layout. +- `dst` element size must be 1, 2, or 4 bytes. + +**Example**: + +```python +pto.tile.gatherb(src_tile, offset_tile, dst_tile) +``` + +The low-level aliases `pto.tsort32`, `pto.tmrgsort`, `pto.tgather`, and +`pto.tgatherb` are also available when a kernel needs to bypass the `pto.tile` +namespace. --- diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 38d24d7bf5..612d7fe5ff 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3605,6 +3605,15 @@ def tgather( ) +def tgatherb(src, offsets, dst): + """``pto.tgatherb`` – tile gather using byte offsets (DPS).""" + _pto.tgatherb( + unwrap_surface_value(src), + unwrap_surface_value(offsets), + unwrap_surface_value(dst), + ) + + def tsel(mask, src0, src1, dst, *, tmp=None): """``pto.tsel ins(mask, src0, src1, tmp) outs(dst)`` with synthesized scratch when omitted.""" resolved_tmp = tmp if tmp is not None else _resolve_selection_tmp(dst, tmp, context="tsel") diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index ca27d42cfb..8f1bd60ea1 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -175,6 +175,7 @@ def rowargmin(src, dst, *, tmp=None): sort32 = staticmethod(_ops.tsort32) mrgsort = staticmethod(_ops.tmrgsort) gather = staticmethod(_ops.tgather) + gatherb = staticmethod(_ops.tgatherb) sel = staticmethod(_ops.tsel) sels = staticmethod(_ops.tsels) diff --git a/ptodsl/ptodsl/tilelib/templates/__init__.py b/ptodsl/ptodsl/tilelib/templates/__init__.py index 6ddd303956..77ad910701 100644 --- a/ptodsl/ptodsl/tilelib/templates/__init__.py +++ b/ptodsl/ptodsl/tilelib/templates/__init__.py @@ -45,6 +45,7 @@ ("a5", "pto.tfillpad"): ".a5.tfillpad", ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand", ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace", + ("a5", "pto.tgatherb"): ".a5.tgatherb", ("a5", "pto.tgemv"): ".a5.tgemv", ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc", ("a5", "pto.tgemv.bias"): ".a5.tgemv_bias", diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py b/ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py new file mode 100644 index 0000000000..7c7da189e0 --- /dev/null +++ b/ptodsl/ptodsl/tilelib/templates/a5/tgatherb.py @@ -0,0 +1,50 @@ +# 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.tgatherb.""" + +from ptodsl import pto +import ptodsl.tilelib as tilelib +from ._common import NUMERIC_DTYPES + + +def gatherb_dtype_signatures(dtypes=NUMERIC_DTYPES): + return [(dtype, "ui32", dtype) for dtype in dtypes] + + +@tilelib.tile_template( + op="pto.tgatherb", + target="a5", + name="template_tgatherb", + dtypes=gatherb_dtype_signatures(), + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False +) +def template_tgatherb(src: pto.Tile, offset: pto.Tile, dst: pto.Tile): + dtype = dst.element_type + dtype_offset = offset.element_type + src_ptr = src.as_ptr() + valid_rows, valid_cols = dst.valid_shape + _, off_valid_cols = offset.valid_shape + lanes = pto.elements_per_vreg(dtype) + blocks_per_vreg = 8 + block_size = lanes // blocks_per_vreg + for row in range(0, valid_rows, 1): + remained = valid_cols + for off_col in range(0, off_valid_cols, blocks_per_vreg): + offset_reg = pto.vlds(offset[row, off_col:]) + mask_dst, _ = pto.make_mask(dtype, remained) + mask, remained = pto.make_mask(dtype_offset, remained) + dst_reg = pto.vgatherb(src_ptr, offset_reg, mask) + dst_col = off_col * block_size + pto.vsts(dst_reg, dst[row, dst_col:], mask_dst) + return diff --git a/test/tilelib-st/a5/tgatherb/case.py b/test/tilelib-st/a5/tgatherb/case.py new file mode 100644 index 0000000000..23f2ae3cdc --- /dev/null +++ b/test/tilelib-st/a5/tgatherb/case.py @@ -0,0 +1,171 @@ +#!/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/tgatherb. +# +# tgatherb gathers elements from a source tile into a destination tile using +# byte offsets. Each element of *offset* is a byte address into the flat byte +# representation of *src*; the element at that byte position is written into +# the corresponding position of *dst*. + +from pathlib import Path +import sys + +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 the numpy dtype corresponding to a pto scalar-dtype name.""" + return PTO_TO_NP_DTYPE[pto_type] + + +# Each case is (name, dtype, src_shape, dst_shape). +# The offset tile has shape [dst_rows, dst_cols // elems_per_block] and contains +# byte offsets into the flat src tile. +# elems_per_block = 32 // dtype.itemsize +CASE_SHAPES = [ + ("f32_2x128", pto.f32, (2, 128), (2, 128)), + ("i32_2x128", pto.i32, (2, 128), (2, 128)), + ("ui32_2x128", pto.ui32, (2, 128), (2, 128)), + ("i16_1x32768", pto.i16, (1, 32768), (1, 32768)), + ("ui16_257x128", pto.ui16, (257, 128), (257, 128)), + ("f16_1x32768", pto.f16, (1, 32768), (1, 32768)), + ("i8_2x256", pto.i8, (2, 256), (2, 256)), + ("i8_2x32768", pto.i8, (2, 32768), (2, 32768)), + ("ui8_2x32768", pto.ui8, (2, 32768), (2, 32768)), +] + + +def _tgatherb_body(src_ptr, offset_ptr, dst_ptr, *, src_rows, src_cols, dst_rows, dst_cols, dtype): + """Shared kernel body for the tgatherb cases. + + Loads *src* and *offset* tiles from GM, performs block-gather using + ``pto.vgatherb``, and stores *dst* back to GM. + """ + block_size_elem = 32 // np.dtype(npy_dtype(dtype)).itemsize + offset_cols = dst_cols // block_size_elem + + src_view = pto.make_tensor_view(src_ptr, shape=[src_rows, src_cols], strides=[src_cols, 1]) + offset_view = pto.make_tensor_view(offset_ptr, shape=[dst_rows, offset_cols], strides=[offset_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) + offset_tile = pto.alloc_tile(shape=[dst_rows, offset_cols], dtype=pto.ui32) + dst_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=dtype) + + pto.tile.load(src_view, src_tile) + pto.tile.load(offset_view, offset_tile) + pto.tile.gatherb(src_tile, offset_tile, dst_tile) + pto.tile.store(dst_tile, dst_view) + + +# One decorated kernel per case, each binding static shapes at definition time. +_tgatherb_kernels = {} +for _name, _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, dtype=_dtype, kernel_name=f"tgatherb_{_name}"): + @pto.jit( + name=kernel_name, + target="a5" + ) + def _kernel( + src_ptr: pto.ptr(dtype, "gm"), + offset_ptr: pto.ptr(pto.ui32, "gm"), + dst_ptr: pto.ptr(dtype, "gm"), + ): + _tgatherb_body( + src_ptr, offset_ptr, dst_ptr, + src_rows=sr, src_cols=sc, dst_rows=dr, dst_cols=dc, dtype=dtype, + ) + + return _kernel + + _tgatherb_kernels[_name] = _make() + + +def _get_data_size(dtype): + """Return the element size in bytes for a given numpy dtype.""" + return np.dtype(dtype).itemsize + + +def _make_inputs(name, dtype, src_shape, dst_shape): + dtype = npy_dtype(dtype) + data_size = _get_data_size(dtype) + block_size_elem = 32 // data_size + + rng_seed = hash(name) & 0xFFFFFFFF + rng = np.random.RandomState(rng_seed) + + src = np.arange(np.prod(src_shape)).astype(dtype) + + offset_col = dst_shape[1] // block_size_elem + offset_shape = (dst_shape[0], offset_col) + offset = np.zeros(np.prod(offset_shape)).astype(np.int32) + for i in range(len(offset)): + offset[i] = i * 32 + offset = offset.reshape(offset_shape) + + return [src, offset] + + +def _make_expected(src, offsets, dtype): + dtype = npy_dtype(dtype) + data_size = _get_data_size(dtype) + block_size_elem = 32 // data_size + + flat_src = src.reshape(-1) + dst = np.zeros(offsets.size * block_size_elem).astype(dtype) + flat_offsets = offsets.ravel() + + count = 0 + for i in range(len(flat_offsets)): + byte_off = int(flat_offsets[i]) + block_start = byte_off // data_size + for j in range(block_size_elem): + dst[count] = flat_src[block_start + j] + count += 1 + + return dst.reshape(src.shape) + + +CASES = [] +for _name, _dtype, _src_shape, _dst_shape in CASE_SHAPES: + CASES.append( + golden_output_case( + "tgatherb_" + _name, + _tgatherb_kernels[_name], + inputs=lambda _n=_name, _d=_dtype, _ss=_src_shape, _ds=_dst_shape: _make_inputs(_n, _d, _ss, _ds), + expected=lambda src, offsets, _d=_dtype: _make_expected(src, offsets, _d), + rtol=1e-6, + atol=1e-6, + ) + ) + + +auto_main(globals()) From 796c9db65d53580e7d3958153c2156829a517fb5 Mon Sep 17 00:00:00 2001 From: Giacomo Castiglioni Date: Tue, 21 Jul 2026 13:42:25 +0200 Subject: [PATCH 2/2] feat(vpto): add A2/A3 gather lowering --- docs/designs/a2a3-vpto-tgather.md | 297 ++++++++++++++++++ .../cce-mlir-vgather-intrinsic-mismatch.md | 142 +++++++++ include/PTO/IR/VPTOUbOps.td | 67 ++++ lib/PTO/IR/PTO.cpp | 27 ++ lib/PTO/IR/VPTO.cpp | 43 +++ .../Transforms/InsertSync/SyncMacroModel.cpp | 13 + lib/PTO/Transforms/LowerPTOToUBufOps.cpp | 196 ++++++++++++ lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 157 +++++++++ ptodsl/tests/e2e/common.py | 232 ++++++++++++++ ptodsl/tests/e2e/test_gather.py | 148 +++++++++ .../tgatherb_a2a3_offsets_dtype_invalid.pto | 20 ++ .../lit/pto/tgatherb_a2a3_offsets_invalid.pto | 20 ++ .../ub/tile_to_ub/gather/tgather_index_a3.pto | 27 ++ .../vpto/ub/tile_to_ub/gather/tgatherb_a3.pto | 25 ++ .../gather/tgatherb_multi_repeat.pto | 25 ++ .../gather/tgather_index_full_chain.pto | 27 ++ .../ub_to_llvm/gather/tgatherb_full_chain.pto | 26 ++ .../ub/ub_to_llvm/gather/vgather_dtypes.pto | 35 +++ .../ub/ub_to_llvm/gather/vgatherb_dtypes.pto | 35 +++ test/lit/vpto/ub/ub_vgather_parse.pto | 26 ++ test/lit/vpto/ub/ub_vgatherb_parse.pto | 29 ++ 21 files changed, 1617 insertions(+) create mode 100644 docs/designs/a2a3-vpto-tgather.md create mode 100644 docs/designs/cce-mlir-vgather-intrinsic-mismatch.md create mode 100644 ptodsl/tests/e2e/test_gather.py create mode 100644 test/lit/pto/tgatherb_a2a3_offsets_dtype_invalid.pto create mode 100644 test/lit/pto/tgatherb_a2a3_offsets_invalid.pto create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/tgather_index_a3.pto create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/tgatherb_a3.pto create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/tgatherb_multi_repeat.pto create mode 100644 test/lit/vpto/ub/ub_to_llvm/gather/tgather_index_full_chain.pto create mode 100644 test/lit/vpto/ub/ub_to_llvm/gather/tgatherb_full_chain.pto create mode 100644 test/lit/vpto/ub/ub_to_llvm/gather/vgather_dtypes.pto create mode 100644 test/lit/vpto/ub/ub_to_llvm/gather/vgatherb_dtypes.pto create mode 100644 test/lit/vpto/ub/ub_vgather_parse.pto create mode 100644 test/lit/vpto/ub/ub_vgatherb_parse.pto diff --git a/docs/designs/a2a3-vpto-tgather.md b/docs/designs/a2a3-vpto-tgather.md new file mode 100644 index 0000000000..24ca540456 --- /dev/null +++ b/docs/designs/a2a3-vpto-tgather.md @@ -0,0 +1,297 @@ +# A2/A3 VPTO Gather Lowering — Design Notes + +Scope: bring tile-level gather (`pto.tgather`, `pto.tgatherb`) to the A2/A3 +`vpto` backend so it reaches feature parity with the EmitC path for gather, +using the flat `pto.ub.*` UB-pointer op layer (the same layer as the +elementwise `ub.vadd` family). The A5 `vreg`/`vgather2` path is **out of +scope** for A2/A3 and must not be used. + +## Background: two lowering tiers for gather + +PTOAS has two parallel lowering tiers for gather, one per backend family: + +| Backend | Gather tier | Primitive | Where | +|---|---|---|---| +| A5 (`vpto`) | vreg | `vgather2` / `vgather2_bc` (vreg operands) | `VPTOLLVMEmitter.cpp:7926` lowers to `llvm.hivm.vgather2.v300.*` | +| A2/A3 (`vpto`) | flat UB pointer (`pto.ub.*`) | raw CCE `vgather` / `vgatherb` (UB pointers) | implemented by this branch | +| Any arch (EmitC) | tile-level CCE call | `TGATHER` / `TGATHERB` opaque calls | `PTOToEmitC.cpp:9772` (tgather), `:9875` (tgatherb) | + +The EmitC path already covers gather for A2/A3 (`ptoas --pto-arch=a3 %s` → +bisheng → pto-isa `TGATHER*`). The gap this design closes is the **vpto/A2/A3** +path, which previously had no gather lowering at all. + +## Raw CCE primitives for A2/A3 gather + +Read from the pto-isa source (`include/pto/npu/a2a3/`). A2/A3 gather is built +on flat `__ubuf__` pointer intrinsics with repeat/stride configuration — the +same shape as the `pto.ub.*` elementwise ops — **not** vreg. + +### `vgather` — element-indexed gather (TGather index/mask/compare form) + +`npu/a2a3/TGather.hpp` implements the index/mask/compare forms on top of the +raw CCE `vgather`: + +```cpp +__ubuf__ T* dstPtr = (__ubuf__ T*)__cce_get_tile_ptr(dst); +__ubuf__ T* src0Ptr = (__ubuf__ T*)__cce_get_tile_ptr(src0); +... +constexpr unsigned elementsPerRepeat = REPEAT_BYTE / sizeof(T); +unsigned numRepeatPerLine = validCol / elementsPerRepeat; +... +vgather((__ubuf__ uint32_t*)(dstPtr + i * TShape1), ...); // 32-bit lanes +vgather((__ubuf__ uint16_t*)(dstPtr + i * TShape1), ...); // 16-bit lanes +``` + +Three forms live in `TGather.hpp`: index (`TGather`), mask +(`TGather`), and compare (`TGather_cmp<...,CmpMode,offset>`). +A2/A3 supports `CmpMode::GT | EQ` only (`TGather.hpp:174`). + +### `vgatherb` — byte-offset gather (TGatherB) + +`npu/a2a3/TGatherB.hpp` wraps the raw CCE `vgatherb`: + +```cpp +// GatherBInstr +vgatherb((__ubuf__ T*)dst, offset /*__ubuf__ uint32_t**/, srcAddr, + dstRepeatStride, /*deCal=*/1, repeats); +// b8 variant widens to uint16_t lanes +vgatherb((__ubuf__ uint16_t*)dst, offset, srcAddr, dstRepeatStride, 1, repeats); +``` + +The tile driver (`GatherBlockHead`/`GatherBlockTail`) loops over +`validRow × numRepeatPerLine` with the `REPEAT_MAX` chunking — structurally +identical to `LowerPTOToUBufOps`'s `dispatch` head/tail repeat loop. + +### What A2/A3 does *not* have + +- **No `MGather`** for A2/A3 — GM gather-load (`MGATHER`) + is A5-only (`npu/a5/MGather.hpp`, implemented via `cce::async_invoke< + simt_mgather_*_kernel>`). A2/A3 has `TGather`, `TGatherB`, `TScatter`, but no + `MGather.hpp`. +- **No vreg** — A5's `vgather2`/`vgather2_bc` (`npu/a5/TGather.hpp:51,79,108`) + must not be used on A2/A3. + +## EmitC reference (the contract to match) + +The EmitC lowering emits arch-agnostic CCE opaque calls that bisheng resolves +to the pto-isa templates above. This is the semantic contract the new vpto +lowering must reproduce: + +| Tile op | EmitC call (`PTOToEmitC.cpp`) | Forms | +|---|---|---| +| `pto.tgather` | `emitc::CallOpaqueOp "TGATHER"` (`:9772`) | index `TGATHER(dst,src0,indices,tmp)`; compare `TGATHER(dst,src0,k,tmp,c,offset)`; mask `TGATHER(dst,src0)` | +| `pto.tgatherb` | `emitc::CallOpaqueOp "TGATHERB"` (`:9875`) | `TGATHERB(dst, src, offsets)` | +| `pto.mgather` | `emitc::CallOpaqueOp "MGATHER"` (`:3323`) | A5-only; out of scope for A2/A3 vpto | + +MGATHER semantics (A5, for reference): `dst[r,:]=table[idx[r],:]` (Coalesce::Row) +or `dst[i,j]=table[idx[i,j]]` (Coalesce::Elem), with `GatherOOB` modes. + +## Current `pto.ub.*` layer + +Defined in `include/PTO/IR/VPTOUbOps.td`. Before this work it was +**elementwise-only**: + +- Binary: `ub.vadd/vsub/vmul/vdiv/vmax/vmin/vand/vor/vaddrelu` +- Unary: `ub.vnot/vabs/vrelu/vexp/vln/vsqrt/vrsqrt` +- Shift/scalar: `ub.vshl/vshr/vmuls/vadds/vmaxs/vmins` +- Misc: `ub.vdup`, `ub.set_mask[_count|_norm]` + +These map to `llvm.hivm.*` for the dav-c220-vec (A2/A3) backend, operating on +flat `PTO_BufferType` (`!pto.ptr<..., ub>`) operands with the repeat / +block-stride / repeat-stride packed into i64 config words +(`cce_aicore_intrinsics_3101.h` layout). This branch adds the corresponding +flat-UB gather ops. + +## Plan + +### 1. New UB gather ops (`include/PTO/IR/VPTOUbOps.td`) + +Add flat-UB-pointer gather ops mirroring `PTO_UBufBinaryOp`/`PTO_UBufUnaryOp` +(PTO_BufferType operands + i64 repeat/stride config), shaped to match the raw +CCE signatures: + +- `pto.ub.vgather` — element-indexed gather (wraps `vgather`). +- `pto.ub.vgatherb` — byte-offset gather (wraps `vgatherb`): + `(dst, src/offsets, srcAddr, repeat, strides…)`. + +Verifier + `getEffects` follow the existing UB op pattern (`lib/PTO/IR/VPTO.cpp` +UBVaddOp-style). C++ builders generated via tablegen. + +### 2. Tile → UB lowering (`lib/PTO/Transforms/LowerPTOToUBufOps.cpp`) + +Add `TGatherOp` / `TGatherBOp` walks to the pass (a2/a3 arch guard already in +place at `:221-224`), following the `dispatch` head/tail repeat loop: +- `pto.tgather` (index form) → `pto.ub.vgather` +- `pto.tgatherb` → `pto.ub.vgatherb` +- Mask/compare forms: decompose to compare+`vgather` (defer; index form first). + +Reuse the existing helpers (`extractTileShapeInfo`, `addPtr`, mask count/norm, +`TileBufAddrOp` for tile→UB ptr). Match the CCE driver's +`validRow × numRepeatPerLine` chunking. + +### 3. UB → LLVM lowering (`lib/PTO/Transforms/VPTOLLVMEmitter.cpp`) + +The `llvm.hivm` callees are resolved (from `docs/designs/a2a3-vector-builtins.md`, +probe target `dav-c220-vec`): + +| UB op | `llvm.hivm` intrinsic | Intrinsic signature | +|---|---|---| +| `pto.ub.vgather` | `llvm.hivm.VGATHER.b16` / `.b32` | `void (ptr addrspace(6) dst, ptr addrspace(6) src, i64 config)` | +| `pto.ub.vgatherb` | `llvm.hivm.VGATHERB.b16` / `.b32` | `void (ptr addrspace(6) dst, ptr addrspace(6) src, i64 config)` | + +Both take **2 UB pointers + 1 packed `i64` config** — the same shape as the +elementwise UB ops. Mirror `LowerUBufBinaryOpPattern` (`VPTOLLVMEmitter.cpp:4734`) +which packs repeat/stride fields into one `i64` via `maskByte`/`shl`/`OrIOp` +and passes the pointer operands directly. Register in the `patterns.add<…>` +block (`:10955`) and declare legality (`:11110`). Keep `VPTOCANN900LLVMEmitter.cpp` +in sync (dormant via the dispatcher's `return false`, but must not rot). + +The CCE builtin forms `vgatherb(dst, src, offsetAddr, dstRepeatStride, +dstBlockStride, repeat)` get packed into `(dst_ptr, src_ptr, config)`. The +elementwise config bit-layout is documented in `VPTOLLVMEmitter.cpp` (e.g. +`LowerUBufShiftOpPattern:4885`); the **vgather/vgatherb config bit-layout is +not documented in CANN** ("No `// ->` packing comment"), so it was decoded +empirically from the bisheng-emitted LLVM IR. The `.b16`/`.b32` suffix selects +16-bit vs 32-bit element lanes. + +### VGATHERB config bit-layout (decoded from bisheng IR) + +Ground truth obtained by probing the real builtin with +`ccec --cce-aicore-arch=dav-c220-vec -xcce -std=c++17 -fenable-matrix \ +-emit-llvm -c -I probe.cpp` (no `-o`; multi-output) then +`llvm-dis` on the device `.bc` (`*-cce-hiipu64-hisilicon-cce-dav-c220-vec.bc`). + +Emitted call: +```llvm +call void @llvm.hivm.VGATHERB.b32(ptr addrspace(6) %dst, ptr addrspace(6) %offset, i64 %config) +``` +- **2nd pointer operand = the offset buffer** (NOT the source data); the source + data base address is packed *into* `config[31:0]`. +- The `i64` config packs the scalar fields as: + - `config[31:0]` — source data address (`offsetAddr`/`srcAddr`) + - `config[39:32]` — `dstRepeatStride` + - `config[47:40]` — `dstBlockStride` + - `config[55:48]` — reserved (0) + - `config[63:56]` — `repeat` +- Verified with a probe using distinctive values + (`srcAddr=0x12345678, dstRepeatStride=0xab, dstBlockStride=0xcd, repeat=0xef`) + which produced `or` chains at exactly the shifts above + (`0xAB<<32`, `0xCD<<40`, `0xEF<<56`). +- The `llvm.hivm.VGATHERB.b16/.b32.cfg.v220` 6-operand overload also exists + (fields passed directly, no packing) but the 3-operand packed form is the one + bisheng emits from the bottom-level `__builtin_cce_vgatherb`. + +Note: `llvm.hivm.VGATHER.b16/.b32` (element gather) uses the same packed +`(dst, src/addr-ptr, i64 config)` shape; its layout will be confirmed the same +way when `pto.ub.vgather` is implemented (Slice 2). The cce-mlir repo's +4-argument `VGATHER(dst, src, addr, repeat)` lowering is a simplified, +bisheng-unverified model (drops strides, no suffix, no bisheng test) — not used. + +### 4. End-to-end tests (`ptodsl/tests/e2e/test_gather.py`) + +PTODSL-authored via `pto.tile.gatherb` (block form) and `pto.tile.gather` +(index form), a3/vpto, numpy goldens, reusing the +`ptodsl/tests/e2e/common.py` harness. Run on NPU manually; there is no CI NPU +job for these today. + +### 5. Lit tests + +- A3 lit pinning `pto.tgather` (index) → `pto.ub.vgather` under + `pto-lower-to-ubuf-ops` (modeled on `test/lit/vpto/ub/tile_to_ub/elementwise/tadd_a2.pto`). +- Replace the current `tgather_index_a3.pto` (which wrongly expects the A5 + `vgather2`) with the `pto.ub.vgather` expectation. + +## Phasing + +1. **Slice 1 — DONE**: `pto.ub.vgatherb` (byte-offset, `vgatherb`): op def + + `tgatherb → ub.vgatherb` tile→UB lowering (with multi-repeat stride from + `TileShapeInfo`) + `ub.vgatherb → llvm.hivm.VGATHERB.b16/.b32` UB→LLVM + lowering. Lit-tested at all three levels (tile→UB, UB→LLVM, full-chain). +2. **Slice 2 — DONE**: `pto.ub.vgather` (element, `vgather`): op def + + UB→LLVM lowering (`→ llvm.hivm.VGATHER.b16/.b32`) + `TGatherOp` index-form + tile→UB lowering via index→byte-offset decomposition (`ub.vmuls` for byte + offsets + `ub.vgather` with `srcAddr` packed as `offsetAddr`). Lit-tested + at all levels. +3. **Slice 3 — PENDING**: mask / compare tgather forms. These require + compile-time index-buffer generation (expanding mask patterns like P0101 to + explicit element indices, or compare-conditional index selection), which is + non-trivial at the PTO IR level. The pto-isa `a2a3/TGather.hpp` handles + these via template-level loops; the MLIR lowering would need to replicate + that (e.g., generating a constant index buffer + using the index-form path). +4. (Out of scope for A2/A3) `mgather` — A5-only; do not port. + +## References + +- A2/A3 raw CCE gather: `pto-isa:npu/a2a3/TGather.hpp` (`vgather`), + `npu/a2a3/TGatherB.hpp` (`vgatherb`), `npu/a2a3/TScatter.hpp`. +- A5 vreg gather (contrast): `pto-isa:npu/a5/TGather.hpp` (`vgather2`). +- EmitC contract: `PTOToEmitC.cpp:9772` (tgather), `:9875` (tgatherb), + `:3323` (mgather). +- UB op layer: `include/PTO/IR/VPTOUbOps.td`; tile→UB dispatch: + `LowerPTOToUBufOps.cpp`; UB→LLVM: `VPTOLLVMEmitter.cpp`. +- e2e harness: `ptodsl/tests/e2e/common.py`. + +## E2E investigation (NPU hardware) + +### Sync fix +InsertSync did not insert MTE2→V sync around TGatherBOp because it lacked a +sync macro model. Fixed by adding `getTGatherBSyncMacroModel` declaring +tgatherb as a single-phase PIPE_V op reading src+offsets and writing dst. +After the fix, InsertSync correctly inserts set_flag/wait_flag between the +DMA loads and the gather. + +### Correct `vgatherb` semantics +The NPU result initially looked wrong because the test used scalar-gather +expectations. CCE documents `vgatherb` as a **32B block gather**: + +- The `src` operand is the UB buffer containing u32 addresses, not scalar + per-element offsets. +- Each repeat reads 8 u32 addresses, so compact offset rows must be padded to + a multiple of 8 entries even when the output row has a partial tail block. +- Each address must be 32B-aligned. +- The instruction copies 8 source blocks of 32B each into dst. + +Therefore, for a 1×8 f32 tile with all-zero addresses, the expected result is +the first 32B block `[10,20,30,40,50,60,70,80]`, not repeated `src[0]`. +Scalar arbitrary element gather belongs on the `vgather` path, not `vgatherb`. + +The LLVM IR matches the bisheng probe exactly (same intrinsic name, same +operand order, same config packing: srcAddr[31:0], dstRepeatStride[39:32], +dstBlockStride[47:40], repeat[63:56]). + +## E2e NPU debugging session (devices 6,7) + +### Key findings +1. **InsertSync fix**: resolved (sync macro model added for TGatherBOp). +2. **vgatherb reads compact block addresses**: repeat=255 crashes if the + compact address buffer is too small, because 255 repeats read 255×8 u32 + addresses and then follow stale address values. +3. **1×8 all-zero output is correct**: repeat=1 copies the first 32B source + block, which appears as sequential f32 values. +4. **Root cause**: the PTOAS e2e harness and part of the tile→UB lowering used + scalar-offset semantics. The lowering must advance the offset pointer by + 8 address entries per repeat and use a compact offset row stride of one + address per 32B output block, padded to 8-entry repeat boundaries. +5. **srcAddr from castptr**: changed from PtrToInt to castptr operand + trace-back (matches pto-isa Tile metadata approach). Both produce the + same i64 value for compile-time addresses. + +### Next step +Use block-gather e2e tests: compact offset tensors with one 32B-aligned byte +address per output 32B block (padded to 8 entries per repeat), and numpy +goldens that copy whole 32B blocks. + +## Final NPU status + +Validated on 910B2 devices 6/7 with `ptodsl/tests/e2e/test_gather.py`: + +- `tgatherb` block gather: f32 shapes `1x8`, `1x64`, `1x96`, `2x64`, + `1x128`, `4x64`. +- `tgatherb` block gather: f16 shapes `1x16`, `1x128`, `1x192`. +- `tgather` scalar index gather: f32/f16 shapes `1x64`, `2x64`. + +The final scalar `tgather` fix was to keep count-mode mask active across the +`vmuls` and `vgather` pair and pass `srcAddr` as `vgather`'s packed +`offsetAddr`, matching `pto/npu/a2a3/TGather.hpp`. Resetting the mask between +`vmuls` and `vgather` made the hardware read too many address entries and +caused UB out-of-bounds. diff --git a/docs/designs/cce-mlir-vgather-intrinsic-mismatch.md b/docs/designs/cce-mlir-vgather-intrinsic-mismatch.md new file mode 100644 index 0000000000..f1b5639372 --- /dev/null +++ b/docs/designs/cce-mlir-vgather-intrinsic-mismatch.md @@ -0,0 +1,142 @@ +# cce-mlir: vgather/vgatherblock emits the wrong LLVM intrinsic form (mismatched vs bisheng) + +Filing-ready summary of a correctness bug in cce-mlir's HIVM→LLVM lowering for +the gather ops. The emitted intrinsic does not match what the `bisheng`/`ccec` +device compiler actually produces/accepts, so it will fail instruction +selection at device-compile time. + +## Summary + +cce-mlir lowers `cce.vgather` / `cce.vgather_block` to a **4-operand, +no-suffix** LLVM intrinsic: + +```cpp +// lib/HIVM/Transforms/LowerHIVMToLLVM.cpp:698-706 +.Case(...) -> emitCall(b, "llvm.hivm.VGATHER", + {o.getDst(), o.getSrc(), o.getAddr(), o.getRepeat()}); +.Case(...) -> emitCall(b, "llvm.hivm.VGATHERBLOCK", + {o.getDst(), o.getSrc(), o.getAddr(), o.getRepeat()}); +``` + +`emitCall` (`LowerHIVMToLLVM.cpp:135-144`) emits the call **verbatim** — it +appends no element-width suffix and performs no config packing. The HIVM ops +themselves are 4-operand `(dst, src, addr, repeat:i8)` +(`include/HIVM/IR/HIVMOpsVector.td:359-383`). + +## What bisheng actually emits (ground truth) + +Probing the real builtins by compiling them with `ccec` and disassembling the +device bitcode shows bisheng emits a **3-operand, packed-config** intrinsic +with an element-width suffix: + +```llvm +declare void @llvm.hivm.VGATHER.b16 (ptr addrspace(6) dst, ptr addrspace(6) src, i64 config) +declare void @llvm.hivm.VGATHER.b32 (ptr addrspace(6) dst, ptr addrspace(6) src, i64 config) +declare void @llvm.hivm.VGATHERB.b16(ptr addrspace(6) dst, ptr addrspace(6) src, i64 config) +declare void @llvm.hivm.VGATHERB.b32(ptr addrspace(6) dst, ptr addrspace(6) src, i64 config) +``` + +(Method: `ccec --cce-aicore-arch=dav-c220-vec -xcce -std=c++17 -fenable-matrix +-emit-llvm -c -I probe.cpp` (no `-o`, multi-output), then `llvm-dis` +on the device `.bc`. Verified directly on `__builtin_cce_vgather` / +`__builtin_cce_vgatherb`. The same 3-operand packed form is recorded in the +PTOAS `docs/designs/a2a3-vector-builtins.md` inventory, which was generated by +the same probe method.) + +For `VGATHERB` the decoded `i64 config` bit-layout is: + +| bits | field | +|---------------|------------------| +| `[31:0]` | source data address (`offsetAddr`/`srcAddr`) | +| `[39:32]` | `dstRepeatStride` | +| `[47:40]` | `dstBlockStride` | +| `[55:48]` | reserved (0) | +| `[63:56]` | `repeat` | + +(2nd pointer operand = the offset buffer; the source data base address is +packed into `config[31:0]`.) + +## Why this is a bug + +1. **Wrong operand count and no suffix.** bisheng's intrinsic is + `(dst, src/offset_ptr, i64 config)` with a `.b16`/`.b32` suffix. cce-mlir + emits `(dst, src, addr, repeat)` with no suffix. Different signature → + bisheng cannot match it in instruction selection + (`Cannot select: intrinsic VGATHER/VGATHERBLOCK`). +2. **Strides are silently dropped.** The CCE-level `vgatherb` raw signature is + `(dst, src, offsetAddr, dstRepeatStride, dstBlockStride, repeat)` (per + `docs/CCE_RAW_INTRINSIC_AUDIT.md:19`), but `LowerToHIVM.cpp:2104-2113` + forwards only `(dst, src, addr, repeat)` — `dstRepeatStride` and + `dstBlockStride` are discarded, and `offsetAddr`/`srcAddr` is never packed + into a config word. So even if the 4-arg form were accepted, the semantics + would be wrong (default strides, lost source-address field). +3. **No bisheng coverage.** There is no `compile_*gather*_to_bisheng.test` + (cf. the existing `compile_vconv_to_bisheng.test`, + `compile_vbrcb_to_bisheng.test`, etc.), and the only vgather test is a + CCE-dialect lit (`test/CCE/ops/select_ubmove.mlir`) that never reaches + bisheng. The audit doc marks vgather/vgatherb "Done" + (`CCE_RAW_INTRINSIC_AUDIT.md:18-19,52`), but that status does not reflect a + bisheng instruction-selection check. + +## Precedent: this is the same class of bug already hit and fixed for VSEL + +`LowerHIVMToLLVM.cpp:717-724` documents exactly this failure mode for `VSEL`: +the earlier 5-operand form `VSEL.(dst,src0,src1,mode,cfg)` caused +`Cannot select: intrinsic VSEL`; it was fixed by probing bisheng and switching +to the real `VSEL..1(dst, src0, src1, i64 cfg)` with the mode packed into +config byte 6. vgather/vgatherblock have not received the same treatment. + +## Suggested fix + +- Emit the real packed intrinsic: `llvm.hivm.VGATHER[B].{b16|b32}(dst_ptr, + offset_ptr, i64 config)` with the element suffix chosen from the element + type, and pack the config word per the layout above. +- Carry `offsetAddr/srcAddr`, `dstRepeatStride`, `dstBlockStride`, and + `repeat` through CCE→HIVM (currently dropped at `LowerToHIVM.cpp:2104-2113`) + so they can be packed. +- Add a `compile_vgather_to_bisheng.test` / `compile_vgatherblock_to_bisheng.test` + (mirroring the existing `compile_*_to_bisheng.test` files) that compiles the + intrinsic through bisheng and asserts the resulting `llvm.hivm.*` call, so + the lowering cannot silently regress against the device compiler. + +## Final PTOAS findings from NPU bring-up + +The packed intrinsic form above is correct, but the semantic model must also +distinguish `vgather` from `vgatherb`: + +- `vgather` is scalar element gather. The address buffer contains byte offsets; + `offsetAddr`/`srcAddr` is packed into `config[31:0]`. pto-isa keeps + count-mode mask active across the address-generation `vmuls` and `vgather`. + Resetting the mask before `vgather` can make hardware read too many address + entries and crash with UB out-of-bounds. +- `vgatherb` is 32B block gather, not scalar gather. Each repeat reads 8 u32 + 32B-aligned block addresses and copies 8 source blocks of 32B each. Compact + offset rows therefore need one i32 address per output 32B block, padded to a + multiple of 8 entries for tail repeats. +- For `vgatherb`, a 1x8 f32 all-zero-address test correctly produces the first + contiguous 32B source block `[src[0]..src[7]]`; it should not be interpreted + as "offsets ignored" or as scalar repeated `src[0]`. + +PTOAS validates these findings with NPU e2e coverage for f32/f16 `tgatherb` +block gather and f32/f16 `tgather` scalar index gather on 910B2. + +## Reproducer (probe) + +```c +// probe_vgatherb.cpp +#include +#include "pto/pto-inst.hpp" +AICORE void probe_vgatherb(__ubuf__ uint32_t *dst, __ubuf__ uint32_t *offset, + uint32_t srcAddr) { + __builtin_cce_vgatherb(dst, offset, srcAddr, (uint16_t)0x00ab, + (uint8_t)0xcd, (uint8_t)0xef); +} +``` +``` +ccec --cce-aicore-arch=dav-c220-vec -xcce -std=c++17 -fenable-matrix \ + -emit-llvm -c -I probe_vgatherb.cpp +llvm-dis probe_vgatherb-cce-hiipu64-hisilicon-cce-dav-c220-vec.bc -o probe.ll +# probe.ll shows: +# call void @llvm.hivm.VGATHERB.b32(ptr addrspace(6) %dst, ptr addrspace(6) %offset, i64 %config) +# with %config = or(srcAddr&0xffffffff, 0xab<<32, 0xcd<<40, 0xef<<56) +``` diff --git a/include/PTO/IR/VPTOUbOps.td b/include/PTO/IR/VPTOUbOps.td index 6de0c51273..97cc483637 100644 --- a/include/PTO/IR/VPTOUbOps.td +++ b/include/PTO/IR/VPTOUbOps.td @@ -192,6 +192,73 @@ def PTO_UBVdupOp : PTO_Op<"ub.vdup", [ }]; } +def PTO_UBVgatherbOp : PTO_Op<"ub.vgatherb", [ + DeclareOpInterfaceMethods + ]> { + let summary = "Raw ubuf byte-offset gather mapping to the CCE vgatherb intrinsic (a2a3)."; + let description = [{ + Lowers to the raw CCE `vgatherb` intrinsic for the a2a3 (dav-c220-vec) + backend. The bishem-emitted intrinsic is the packed 3-operand form + `llvm.hivm.VGATHERB.b16/.b32(ptr dst, ptr offset, i64 config)`, where the + 2nd pointer is the offset buffer and the source data base address is packed + into `config[31:0]`. Config layout (decoded from bisheng IR — see + docs/designs/a2a3-vpto-tgather.md): srcAddr[31:0], dstRepeatStride[39:32], + dstBlockStride[47:40], reserved[55:48]=0, repeat[63:56]. This is the A2/A3 + UB-pointer gather primitive used to lower `pto.tgatherb`; do not confuse + with the A5 vreg `vgather2`. + }]; + + let arguments = (ins + PTO_BufferType:$dst, + PTO_BufferType:$offset, + PTO_BufferType:$src, + I64:$dstRepeatStride, + I64:$dstBlockStride, + I64:$repeat + ); + + let results = (outs); + + let hasVerifier = 1; + + let assemblyFormat = [{ + $dst `,` $offset `,` $src `,` + $dstRepeatStride `,` $dstBlockStride `,` $repeat + attr-dict `:` type($dst) `,` type($offset) `,` type($src) `,` + type($dstRepeatStride) `,` type($dstBlockStride) `,` type($repeat) + }]; +} + +def PTO_UBVgatherOp : PTO_Op<"ub.vgather", [ + DeclareOpInterfaceMethods + ]> { + let summary = "Raw ubuf element gather mapping to the CCE vgather intrinsic (a2a3)."; + let description = [{ + Lowers to `llvm.hivm.VGATHER.b16/.b32(dst_ptr, src_ptr, i64 config)`. + The 2nd pointer is the address buffer; `offsetAddr` is a base offset added + to each address in the buffer. Config layout: offsetAddr[31:0], + dstRepeatStride[39:32], repeat[63:56]. + }]; + + let arguments = (ins + PTO_BufferType:$dst, + PTO_BufferType:$src, + I64:$offsetAddr, + I64:$dstRepeatStride, + I64:$repeat + ); + + let results = (outs); + + let hasVerifier = 1; + + let assemblyFormat = [{ + $dst `,` $src `,` $offsetAddr `,` $dstRepeatStride `,` $repeat + attr-dict `:` type($dst) `,` type($src) `,` type($offsetAddr) `,` + type($dstRepeatStride) `,` type($repeat) + }]; +} + def PTO_UBSetMaskOp : PTO_Op<"ub.set_mask", [ DeclareOpInterfaceMethods ]> { diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 8034fecb64..ca07f9a4f9 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7726,13 +7726,40 @@ mlir::LogicalResult mlir::pto::TGatherBOp::verify() { FailureOr> elems = verifyCommon(); if (failed(elems)) return failure(); + Type srcTy = getSrc().getType(); + Type offTy = getOffsets().getType(); Type dstTy = getDst().getType(); Type dstElemTy = elems->second; + if (failed(verifyTileBufSameValidShape(*this, srcTy, dstTy, "src", "dst"))) + return failure(); if (!isRowMajorTileBuf(dstTy)) return emitOpError() << "expects dst to use row-major layout"; auto dstBytes = getElemBytes(dstElemTy); if (!dstBytes || (*dstBytes != 1 && *dstBytes != 2 && *dstBytes != 4)) return emitOpError() << "expects dst element size to be 1, 2, or 4 bytes"; + Type offElemTy = getElemTy(offTy); + if (!offElemTy.isInteger(32)) + return emitOpError() << "expects offsets element type to be i32"; + + auto dstValid = getValidShapeVec(dstTy); + auto offValid = getValidShapeVec(offTy); + if (dstValid.size() != 2 || offValid.size() != 2) + return emitOpError() << "expects rank-2 src/offsets/dst tile buffers"; + if (dstValid[0] != ShapedType::kDynamic && + offValid[0] != ShapedType::kDynamic && dstValid[0] != offValid[0]) + return emitOpError() << "expects offsets valid rows to match dst valid rows"; + if (dstValid[1] != ShapedType::kDynamic && + offValid[1] != ShapedType::kDynamic) { + int64_t blockElems = 32 / *dstBytes; + int64_t blocks = (dstValid[1] + blockElems - 1) / blockElems; + int64_t expectedOffsetCols = ((blocks + 7) / 8) * 8; + if (offValid[1] != expectedOffsetCols) { + return emitOpError() + << "expects offsets valid cols to be compact 32B block address " + "count padded to 8 entries; expected " + << expectedOffsetCols << ", got " << offValid[1]; + } + } return mlir::success(); }; diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index c62247f21b..9b58f5486d 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -8383,6 +8383,49 @@ LogicalResult UBVdupOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// UBVgatherbOp +//===----------------------------------------------------------------------===// + +void UBVgatherbOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable()); + effects.emplace_back(MemoryEffects::Read::get(), &getOffsetMutable()); + effects.emplace_back(MemoryEffects::Write::get(), &getDstMutable()); +} + +LogicalResult UBVgatherbOp::verify() { + if (!isBufferLike(getDst().getType()) || !isBufferLike(getOffset().getType()) || + !isBufferLike(getSrc().getType())) + return emitOpError("requires pointer-like operands"); + if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || + classifyMemoryRole(getOffset().getType()) != MemoryRole::UB || + classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) + return emitOpError("requires UB-backed operands"); + return success(); +} + +//===----------------------------------------------------------------------===// +// UBVgatherOp +//===----------------------------------------------------------------------===// + +void UBVgatherOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable()); + effects.emplace_back(MemoryEffects::Write::get(), &getDstMutable()); +} + +LogicalResult UBVgatherOp::verify() { + if (!isBufferLike(getDst().getType()) || !isBufferLike(getSrc().getType())) + return emitOpError("requires pointer-like operands"); + if (classifyMemoryRole(getDst().getType()) != MemoryRole::UB || + classifyMemoryRole(getSrc().getType()) != MemoryRole::UB) + return emitOpError("requires UB-backed operands"); + return success(); +} + //===----------------------------------------------------------------------===// // UBVshlOp //===----------------------------------------------------------------------===// diff --git a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp index 14ecc4d453..1f7ba7ff22 100644 --- a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp @@ -409,6 +409,17 @@ std::optional getMGatherSyncMacroModel(pto::MGatherOp op) { return model; } +// TGatherBOp (byte-offset gather): single-phase PIPE_V op that reads src + +// offsets and writes dst. Explicit model ensures InsertSync detects the +// cross-pipe MTE2→V dependency (DMA load of src/offsets must complete before +// the gather reads them). +std::optional getTGatherBSyncMacroModel(pto::TGatherBOp op) { + SyncMacroModel model; + addPhase(model, PipelineType::PIPE_V, ValueRange{op.getDst()}, + ValueRange{op.getSrc(), op.getOffsets()}); + return model; +} + } // namespace std::optional mlir::pto::getSyncMacroModel(Operation *op) { @@ -420,6 +431,8 @@ std::optional mlir::pto::getSyncMacroModel(Operation *op) { return getTScatterSyncMacroModel(tscatter); if (auto tgather = dyn_cast(op)) return getTGatherSyncMacroModel(tgather); + if (auto tgatherb = dyn_cast(op)) + return getTGatherBSyncMacroModel(tgatherb); if (auto mgather = dyn_cast(op)) return getMGatherSyncMacroModel(mgather); return std::nullopt; diff --git a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp index f9fe13866d..cf0e07d5c9 100644 --- a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp +++ b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp @@ -811,6 +811,202 @@ struct LowerPTOToUBufOpsPass op.erase(); } + // ---- tgatherb → pto.ub.vgatherb (GatherBlockHead/Tail tiling) ---- + // Mirrors the pto-isa a2a3/TGatherB.hpp GatherBlockHead/Tail driver. + // The CCE vgatherb hardware requires specific repeat counts and pointer + // advancement per call. + { + SmallVector ops; + func.walk([&](pto::TGatherBOp op) { ops.push_back(op); }); + for (auto op : ops) { + if (!canLower(op, tileShapes)) + continue; + auto info = extractTileShapeInfo(op, tileShapes); + if (!info) + continue; + + unsigned bse = info->blockSizeElem; + unsigned epr = info->elementsPerRepeat; + int64_t validRow = info->vRows; + int64_t validCol = info->vCols; + if (bse == 0 || epr == 0 || validCol == 0) + continue; + + // GatherBlockHead/Tail parameters (from pto-isa a2a3/TGatherB.hpp). + // vgatherb reads 8 u32 block addresses per repeat. Each address must + // point at a 32B source block; the instruction copies those 8 blocks. + int64_t numRepeatPerLine = validCol / epr; + int64_t numRemainPerLine = validCol % epr; + constexpr int64_t REPEAT_MAX = 255; + constexpr int64_t ADDRS_PER_REPEAT = 8; + + // Compact offset tile layout: one i32 address per 32B output block, + // padded to 8 entries because even a tail repeat reads 8 addresses. + int64_t blocksPerRow = (validCol + bse - 1) / bse; + int64_t offsetStridePerRow = + ((blocksPerRow + ADDRS_PER_REPEAT - 1) / ADDRS_PER_REPEAT) * + ADDRS_PER_REPEAT; + + Location loc = op.getLoc(); + builder.setInsertionPoint(op); + + Type dstElemTy = getStoredElemType(op.getDst().getType()); + if (!dstElemTy) + continue; + auto dstPtrType = getUBPtrType(ctx, dstElemTy); + auto offPtrType = getUBPtrType(ctx, builder.getI32Type()); + + auto emitAddr = [&](Value tile, pto::PtrType ty) -> Value { + if (isa(tile.getType())) + return tile; + return builder.create(loc, ty, tile).getDst(); + }; + + Value dstBase = emitAddr(op.getDst(), dstPtrType); + Value offBase = emitAddr(op.getOffsets(), offPtrType); + Value srcBase = emitAddr(op.getSrc(), dstPtrType); + + auto emitGatherb = [&](Value dst, Value off, int64_t repStride, + int64_t repeat) { + builder.create( + loc, dst, off, srcBase, i64c(repStride, loc, builder), + i64c(1, loc, builder), i64c(repeat, loc, builder)); + }; + + // ---- GatherBlockHead: process full epr-element blocks ---- + if (numRepeatPerLine > 0) { + int64_t numLoop = numRepeatPerLine / REPEAT_MAX; + int64_t remainAfterLoop = numRepeatPerLine % REPEAT_MAX; + + for (int64_t i = 0; i < validRow; ++i) { + int64_t rowOff = i * validCol; + int64_t offRowOff = i * offsetStridePerRow; + + for (int64_t j = 0; j < numLoop; ++j) { + int64_t elemOff = rowOff + j * epr * REPEAT_MAX; + int64_t offOff = offRowOff + j * ADDRS_PER_REPEAT * REPEAT_MAX; + Value dstAdv = addPtr(loc, builder, dstBase, dstPtrType, + idxc(elemOff, loc, builder)); + Value offAdv = addPtr(loc, builder, offBase, offPtrType, + idxc(offOff, loc, builder)); + emitGatherb(dstAdv, offAdv, ADDRS_PER_REPEAT, REPEAT_MAX); + } + if (remainAfterLoop > 0) { + int64_t elemOff = rowOff + numLoop * epr * REPEAT_MAX; + int64_t offOff = + offRowOff + numLoop * ADDRS_PER_REPEAT * REPEAT_MAX; + Value dstAdv = addPtr(loc, builder, dstBase, dstPtrType, + idxc(elemOff, loc, builder)); + Value offAdv = addPtr(loc, builder, offBase, offPtrType, + idxc(offOff, loc, builder)); + emitGatherb(dstAdv, offAdv, ADDRS_PER_REPEAT, remainAfterLoop); + } + } + } + + // ---- GatherBlockTail: process remaining elements ---- + if (numRemainPerLine > 0) { + int64_t tailElemOff = numRepeatPerLine * epr; + int64_t tailRepStride = validCol / bse; + if (tailRepStride == 0) + tailRepStride = 1; + + for (int64_t i = 0; i < validRow; ++i) { + int64_t elemOff = i * validCol + tailElemOff; + int64_t offOff = i * offsetStridePerRow + tailElemOff / bse; + Value dstAdv = addPtr(loc, builder, dstBase, dstPtrType, + idxc(elemOff, loc, builder)); + Value offAdv = addPtr(loc, builder, offBase, offPtrType, + idxc(offOff, loc, builder)); + emitGatherb(dstAdv, offAdv, tailRepStride, 1); + } + } + + op.erase(); + } + } + + // ---- tgather (index form) → ub.vmuls + ub.vgather ---- + // Decomposes element-index gather into byte offsets (vmuls), then passes + // the source tile address as vgather's offsetAddr config. This mirrors + // pto-isa a2a3/TGather.hpp. + { + SmallVector ops; + func.walk([&](pto::TGatherOp op) { ops.push_back(op); }); + for (auto op : ops) { + if (!op.hasIndexForm()) + continue; + if (!canLower(op, tileShapes)) + continue; + auto info = extractTileShapeInfo(op, tileShapes); + if (!info) + continue; + int64_t epr = info->elementsPerRepeat; + if (epr == 0) + continue; + + // Extract the source base address (i64) from the CastPtrOp that + // defines the src tile pointer. Only handle alloc_tile-backed src. + Value srcAddr; + if (auto *defOp = op.getSrc().getDefiningOp()) + if (isa(defOp)) + srcAddr = defOp->getOperand(0); + if (!srcAddr) + continue; + + Location loc = op.getLoc(); + builder.setInsertionPoint(op); + + auto i32PtrType = getUBPtrType(ctx, builder.getI32Type()); + auto emitAddr = [&](Value tile, pto::PtrType ty) -> Value { + if (isa(tile.getType())) + return tile; + return builder.create(loc, ty, tile).getDst(); + }; + + Value indicesPtr = emitAddr(op.getIndices(), i32PtrType); + Value tmpPtr = emitAddr(op.getTmp(), i32PtrType); + + Type dstElemTy = getStoredElemType(op.getDst().getType()); + if (!dstElemTy) + continue; + unsigned elemSize = getElementSize(dstElemTy); + if (elemSize == 0) + continue; + auto dstPtrType = getUBPtrType(ctx, dstElemTy); + Value dstPtr = emitAddr(op.getDst(), dstPtrType); + + auto pipeV = pto::PipeAttr::get(ctx, pto::PIPE::PIPE_V); + for (int64_t i = 0; i < info->vRows; ++i) { + for (int64_t col = 0; col < info->vCols; col += epr) { + int64_t chunkElems = std::min(epr, info->vCols - col); + Value rowOff = idxc(i * info->cols + col, loc, builder); + Value tmpRow = addPtr(loc, builder, tmpPtr, i32PtrType, rowOff); + Value idxRow = addPtr(loc, builder, indicesPtr, i32PtrType, rowOff); + Value dstRow = addPtr(loc, builder, dstPtr, dstPtrType, rowOff); + + builder.create(loc); + builder.create(loc, + i64c(chunkElems, loc, builder), + i64c0(loc, builder)); + // tmp = indices * elemSize (byte offsets). Keep count-mode mask + // active for vgather, matching pto-isa a2a3/TGather.hpp. + builder.create( + loc, tmpRow, idxRow, i64c(elemSize, loc, builder), + i64c1(loc, builder), i64c1(loc, builder), i64c1(loc, builder), + i64c8(loc, builder), i64c8(loc, builder)); + builder.create(loc, pipeV); + builder.create( + loc, dstRow, tmpRow, srcAddr, i64c(8, loc, builder), + i64c1(loc, builder)); + } + } + builder.create(loc); + fullMask(loc, builder); + op.erase(); + } + } + // ---- cleanup dead PTO ops ---- SmallVector toErase; func.walk([&](Operation *op) { diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 3ba307ffa2..bd4370cd43 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -4867,6 +4867,157 @@ class LowerUBufBinaryOpPattern final : public OpConversionPattern { LoweringState &state; }; +// pto.ub.vgatherb -> llvm.hivm.VGATHERB.b16/.b32(dst_ptr, offset_ptr, i64 config) +// Config (decoded from bisheng IR, see docs/designs/a2a3-vpto-tgather.md): +// srcAddr[31:0] | dstRepeatStride[39:32] | dstBlockStride[47:40] +// | reserved[55:48]=0 | repeat[63:56] +// The 2nd pointer operand is the offset buffer; the source data base address +// (low 32 bits of the src pointer) is packed into config[31:0]. +class LowerUBVgatherbOpPattern final + : public OpConversionPattern { +public: + explicit LowerUBVgatherbOpPattern(TypeConverter &typeConverter, + MLIRContext *context, LoweringState &state) + : OpConversionPattern(typeConverter, context), + state(state) {} + + LogicalResult + matchAndRewrite(pto::UBVgatherbOp op, pto::UBVgatherbOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Value dst = adaptor.getDst(); + Value offset = adaptor.getOffset(); + Value src = adaptor.getSrc(); + if (!dst || !offset || !src || + !isa(dst.getType()) || + !isa(offset.getType()) || + !isa(src.getType())) + return rewriter.notifyMatchFailure( + op, "unexpected converted ub.vgatherb operand types"); + + auto ptrType = mlir::cast(op.getDst().getType()); + Type elemType = ptrType.getElementType(); + unsigned width = pto::getPTOStorageElemBitWidth(elemType); + if (width != 16 && width != 32) + return rewriter.notifyMatchFailure( + op, "unsupported element width for ub.vgatherb"); + std::string calleeName = + std::string("llvm.hivm.VGATHERB.") + ((width == 16) ? "b16" : "b32"); + + Location loc = op.getLoc(); + auto i64Ty = rewriter.getI64Type(); + auto constI64 = [&](uint64_t v) -> Value { + return rewriter.create(loc, + rewriter.getI64IntegerAttr(v)); + }; + auto getI64 = [&](Value v) -> Value { + return castIntegerLikeTo(op, v, i64Ty); + }; + auto maskByte = [&](Value v) -> Value { + return rewriter.create(loc, v, constI64(0xff)); + }; + auto shl = [&](Value v, uint64_t amount) -> Value { + return rewriter.create(loc, v, constI64(amount)); + }; + + // config[31:0] = source data address (low 32 bits of the src pointer). + // Trace back through castptr to get the planned UB offset, matching the + // address loaded from Tile host_ptr metadata by the PTO-ISA reference. + Value srcAddr; + if (auto *defOp = op.getSrc().getDefiningOp()) { + if (auto castOp = dyn_cast(defOp)) + srcAddr = castOp.getOperand(); + } + if (!srcAddr) + srcAddr = rewriter.create(loc, i64Ty, src); + Value config = + rewriter.create(loc, srcAddr, constI64(0xffffffff)); + config = rewriter.create( + loc, config, shl(maskByte(getI64(adaptor.getDstRepeatStride())), 32)); + config = rewriter.create( + loc, config, shl(maskByte(getI64(adaptor.getDstBlockStride())), 40)); + config = rewriter.create( + loc, config, shl(maskByte(getI64(adaptor.getRepeat())), 56)); + + auto funcType = rewriter.getFunctionType( + TypeRange{dst.getType(), offset.getType(), rewriter.getI64Type()}, + TypeRange{}); + rewriter.create(op.getLoc(), calleeName, TypeRange{}, + ValueRange{dst, offset, config}); + state.plannedDecls.push_back(PlannedDecl{calleeName, funcType}); + rewriter.eraseOp(op); + return success(); + } + +private: + LoweringState &state; +}; + +// pto.ub.vgather -> llvm.hivm.VGATHER.b16/.b32(dst_ptr, src_ptr, i64 config) +// Config: offsetAddr[31:0] | dstRepeatStride[39:32] | repeat[63:56]. +class LowerUBVgatherOpPattern final + : public OpConversionPattern { +public: + explicit LowerUBVgatherOpPattern(TypeConverter &typeConverter, + MLIRContext *context, LoweringState &state) + : OpConversionPattern(typeConverter, context), + state(state) {} + + LogicalResult + matchAndRewrite(pto::UBVgatherOp op, pto::UBVgatherOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Value dst = adaptor.getDst(); + Value src = adaptor.getSrc(); + if (!dst || !src || !isa(dst.getType()) || + !isa(src.getType())) + return rewriter.notifyMatchFailure( + op, "unexpected converted ub.vgather operand types"); + + auto ptrType = mlir::cast(op.getDst().getType()); + Type elemType = ptrType.getElementType(); + unsigned width = pto::getPTOStorageElemBitWidth(elemType); + if (width != 16 && width != 32) + return rewriter.notifyMatchFailure( + op, "unsupported element width for ub.vgather"); + std::string calleeName = + std::string("llvm.hivm.VGATHER.") + ((width == 16) ? "b16" : "b32"); + + Location loc = op.getLoc(); + auto i64Ty = rewriter.getI64Type(); + auto constI64 = [&](uint64_t v) -> Value { + return rewriter.create(loc, + rewriter.getI64IntegerAttr(v)); + }; + auto getI64 = [&](Value v) -> Value { + return castIntegerLikeTo(op, v, i64Ty); + }; + auto maskByte = [&](Value v) -> Value { + return rewriter.create(loc, v, constI64(0xff)); + }; + auto shl = [&](Value v, uint64_t amount) -> Value { + return rewriter.create(loc, v, constI64(amount)); + }; + + Value config = rewriter.create( + loc, getI64(adaptor.getOffsetAddr()), constI64(0xffffffff)); + config = rewriter.create( + loc, config, shl(maskByte(getI64(adaptor.getDstRepeatStride())), 32)); + config = rewriter.create( + loc, config, shl(maskByte(getI64(adaptor.getRepeat())), 56)); + + auto funcType = rewriter.getFunctionType( + TypeRange{dst.getType(), src.getType(), rewriter.getI64Type()}, + TypeRange{}); + rewriter.create(op.getLoc(), calleeName, TypeRange{}, + ValueRange{dst, src, config}); + state.plannedDecls.push_back(PlannedDecl{calleeName, funcType}); + rewriter.eraseOp(op); + return success(); + } + +private: + LoweringState &state; +}; + template class LowerUBufShiftOpPattern final : public OpConversionPattern { public: @@ -11190,6 +11341,10 @@ static void populateVPTOOpLoweringPatterns(VPTOTypeConverter &typeConverter, typeConverter, patterns.getContext(), state); patterns.add( typeConverter, patterns.getContext(), state); + patterns.add( + typeConverter, patterns.getContext(), state); + patterns.add( + typeConverter, patterns.getContext(), state); patterns.add( typeConverter, patterns.getContext(), state); patterns.add( @@ -11340,6 +11495,8 @@ static void configureVPTOOpLoweringTarget(ConversionTarget &target, target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); + target.addIllegalOp(); + target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); target.addIllegalOp(); diff --git a/ptodsl/tests/e2e/common.py b/ptodsl/tests/e2e/common.py index 7db6c9bf10..d8a8096539 100644 --- a/ptodsl/tests/e2e/common.py +++ b/ptodsl/tests/e2e/common.py @@ -484,3 +484,235 @@ def launch_and_check_scalar( actual = z.cpu().numpy() np.testing.assert_allclose(actual, ref, rtol=rtol, atol=atol) return compile_s, launch_s + + +# --------------------------------------------------------------------------- +# Gather e2e helpers +# --------------------------------------------------------------------------- + +def make_gatherb_kernel( + rows: int, + cols: int, + dtype_str: str = "float32", + target: str = "a3", + backend: str = "vpto", + kernel_kind: str = "vector", +): + """Return a ``@pto.jit`` KernelHandle for pto.tile.gatherb (block gather). + + Kernel signature: (Src_ptr f32 GM, Off_ptr i32 GM, Dst_ptr f32 GM). + Loads src and compact 32B block addresses into UB tiles, calls gatherb, + and stores the gathered blocks. + """ + pto_dtype = getattr(pto, dtype_str) + off_dtype = pto.i32 + fn_name = f"gatherb_{dtype_str}_{rows}x{cols}" + + def kernel_body( + Src_ptr: pto.ptr(pto_dtype, "gm"), + Off_ptr: pto.ptr(off_dtype, "gm"), + Dst_ptr: pto.ptr(pto_dtype, "gm"), + ) -> None: + c0 = pto.const(0) + c1 = pto.const(1) + c_rows = pto.const(rows) + c_cols = pto.const(cols) + c_elems = pto.const(rows * cols) + block_elems = 8 if dtype_str == "float32" else 16 + block_cols = (cols + block_elems - 1) // block_elems + block_cols = ((block_cols + 7) // 8) * 8 + block_total = rows * block_cols + c_block_cols = pto.const(block_cols) + c_block_total = pto.const(block_total) + + shape = [c1, c1, c1, c_rows, c_cols] + strides = [c_elems, c_elems, c_elems, c_cols, c1] + off_shape = [c1, c1, c1, c_rows, c_block_cols] + off_strides = [c_block_total, c_block_total, c_block_total, + c_block_cols, c1] + off = [c0, c0, c0, c0, c0] + + src_view = pto.make_tensor_view(Src_ptr, shape=shape, strides=strides) + off_view = pto.make_tensor_view(Off_ptr, shape=off_shape, + strides=off_strides) + dst_view = pto.make_tensor_view(Dst_ptr, shape=shape, strides=strides) + + src_part = pto.partition_view(src_view, offsets=off, sizes=shape) + off_part = pto.partition_view(off_view, offsets=off, sizes=off_shape) + dst_part = pto.partition_view(dst_view, offsets=off, sizes=shape) + + src_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto_dtype) + off_tile = pto.alloc_tile(shape=[rows, block_cols], dtype=off_dtype) + dst_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto_dtype) + + pto.tile.load(src_part, src_tile) + pto.tile.load(off_part, off_tile) + pto.tile.gatherb(src_tile, off_tile, dst_tile) + pto.tile.store(dst_tile, dst_part) + + kernel_body.__name__ = fn_name + return pto.jit( + name=fn_name, + kernel_kind=kernel_kind, + target=target, + backend=backend, + )(kernel_body) + + +def launch_and_check_gather( + *, + rows: int, + cols: int, + dtype_str: str, + torch, + kernel_handle, + seed: int = 42, + rtol: float = 1e-6, + atol: float = 1e-6, +): + """Compile, launch, and numerical-check a gatherb kernel. + + ``vgatherb`` is a 32B block gather: each i32 offset entry selects one + 32B source block, and the instruction copies that whole block to dst. + """ + rng = np.random.RandomState(seed) + total = rows * cols + elem_size = 4 if dtype_str == "float32" else 2 + block_elems = 32 // elem_size + valid_block_cols = (cols + block_elems - 1) // block_elems + block_cols = ((valid_block_cols + 7) // 8) * 8 + valid_blocks = rows * valid_block_cols + total_blocks = rows * block_cols + np_dtype = np.float32 if dtype_str == "float32" else np.float16 + + # Random source data (small integers for exact results). + src_flat = rng.randint(1, 10, size=(total,)).astype(np_dtype) + # Random valid 32B-aligned byte offsets, one per output 32B block. + valid_src_blocks = valid_blocks + block_idx = rng.randint(0, valid_src_blocks, size=(total_blocks,)) + off_flat = (block_idx * 32).astype(np.int32) + + golden_flat = np.empty((total,), dtype=np_dtype) + for row in range(rows): + for col_block in range(valid_block_cols): + out_block = row * valid_block_cols + col_block + src_block = block_idx[row * block_cols + col_block] + out_start = out_block * block_elems + src_start = src_block * block_elems + count = min(block_elems, total - out_start, total - src_start) + golden_flat[out_start:out_start + count] = src_flat[src_start:src_start + count] + golden = golden_flat.reshape(rows, cols) + + torch_dt = _torch_dtype(torch, dtype_str) + src_dev = torch.from_numpy(src_flat).to(device="npu:0", dtype=torch_dt) + off_dev = torch.from_numpy(off_flat).to(device="npu:0", dtype=torch.int32) + dst_dev = torch.empty(total, dtype=torch_dt, device="npu:0") + stream = _npu_stream(torch) + + t0 = time.perf_counter() + compiled = kernel_handle.compile() + compile_s = time.perf_counter() - t0 + + t0 = time.perf_counter() + compiled[1, stream](src_dev.data_ptr(), off_dev.data_ptr(), dst_dev.data_ptr()) + torch.npu.synchronize() + launch_s = time.perf_counter() - t0 + + actual = dst_dev.cpu().numpy().reshape(rows, cols) + np.testing.assert_allclose(actual, golden, rtol=rtol, atol=atol) + return compile_s, launch_s + + +def make_gather_index_kernel( + rows: int, + cols: int, + dtype_str: str = "float32", + target: str = "a3", + backend: str = "vpto", + kernel_kind: str = "vector", +): + """Return a ``@pto.jit`` KernelHandle for scalar index ``pto.tile.gather``.""" + pto_dtype = getattr(pto, dtype_str) + idx_dtype = pto.i32 + fn_name = f"gather_index_{dtype_str}_{rows}x{cols}" + + def kernel_body( + Src_ptr: pto.ptr(pto_dtype, "gm"), + Idx_ptr: pto.ptr(idx_dtype, "gm"), + Dst_ptr: pto.ptr(pto_dtype, "gm"), + ) -> None: + c0 = pto.const(0) + c1 = pto.const(1) + c_rows = pto.const(rows) + c_cols = pto.const(cols) + c_elems = pto.const(rows * cols) + + shape = [c1, c1, c1, c_rows, c_cols] + strides = [c_elems, c_elems, c_elems, c_cols, c1] + off = [c0, c0, c0, c0, c0] + + src_view = pto.make_tensor_view(Src_ptr, shape=shape, strides=strides) + idx_view = pto.make_tensor_view(Idx_ptr, shape=shape, strides=strides) + dst_view = pto.make_tensor_view(Dst_ptr, shape=shape, strides=strides) + + src_part = pto.partition_view(src_view, offsets=off, sizes=shape) + idx_part = pto.partition_view(idx_view, offsets=off, sizes=shape) + dst_part = pto.partition_view(dst_view, offsets=off, sizes=shape) + + src_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto_dtype) + idx_tile = pto.alloc_tile(shape=[rows, cols], dtype=idx_dtype) + tmp_tile = pto.alloc_tile(shape=[rows, cols], dtype=idx_dtype) + dst_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto_dtype) + + pto.tile.load(src_part, src_tile) + pto.tile.load(idx_part, idx_tile) + pto.tile.gather(src_tile, dst_tile, indices=idx_tile, tmp=tmp_tile) + pto.tile.store(dst_tile, dst_part) + + kernel_body.__name__ = fn_name + return pto.jit( + name=fn_name, + kernel_kind=kernel_kind, + target=target, + backend=backend, + )(kernel_body) + + +def launch_and_check_gather_index( + *, + rows: int, + cols: int, + dtype_str: str, + torch, + kernel_handle, + seed: int = 42, + rtol: float = 1e-6, + atol: float = 1e-6, +): + """Compile, launch, and numerical-check scalar index ``pto.tile.gather``.""" + rng = np.random.RandomState(seed) + total = rows * cols + np_dtype = np.float32 if dtype_str == "float32" else np.float16 + + src_flat = rng.randint(1, 10, size=(total,)).astype(np_dtype) + idx_flat = rng.randint(0, total, size=(total,)).astype(np.int32) + golden = src_flat[idx_flat].reshape(rows, cols) + + torch_dt = _torch_dtype(torch, dtype_str) + src_dev = torch.from_numpy(src_flat).to(device="npu:0", dtype=torch_dt) + idx_dev = torch.from_numpy(idx_flat).to(device="npu:0", dtype=torch.int32) + dst_dev = torch.empty(total, dtype=torch_dt, device="npu:0") + stream = _npu_stream(torch) + + t0 = time.perf_counter() + compiled = kernel_handle.compile() + compile_s = time.perf_counter() - t0 + + t0 = time.perf_counter() + compiled[1, stream](src_dev.data_ptr(), idx_dev.data_ptr(), dst_dev.data_ptr()) + torch.npu.synchronize() + launch_s = time.perf_counter() - t0 + + actual = dst_dev.cpu().numpy().reshape(rows, cols) + np.testing.assert_allclose(actual, golden, rtol=rtol, atol=atol) + return compile_s, launch_s diff --git a/ptodsl/tests/e2e/test_gather.py b/ptodsl/tests/e2e/test_gather.py new file mode 100644 index 0000000000..2dc54525ce --- /dev/null +++ b/ptodsl/tests/e2e/test_gather.py @@ -0,0 +1,148 @@ +# 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. + +""" +e2e tests for A3 VPTO gather. + +Exercises the full lowering chain on NPU hardware: + pto.tile.gatherb → pto.tgatherb → pto.ub.vgatherb → llvm.hivm.VGATHERB.b32 + pto.tile.gather → pto.tgather → pto.ub.vgather → llvm.hivm.VGATHER.b32 + +Run: + pytest ptodsl/tests/e2e/test_gather.py -v + +Filter: + pytest ptodsl/tests/e2e/test_gather.py -v -k "gatherb-float32-1x64" +""" + +from __future__ import annotations + +import pytest + +from .common import ( + make_gather_index_kernel, + make_gatherb_kernel, + launch_and_check_gather, + launch_and_check_gather_index, +) + + +# --------------------------------------------------------------------------- +# Shape matrix — f32 (elementsPerRepeat = 64) and f16 +# (elementsPerRepeat = 128). gatherb works in 32B blocks, so the smallest +# useful cases are 8 f32 or 16 f16 values. +# --------------------------------------------------------------------------- + +F32_EPR = 64 +F16_EPR = 128 + +GATHERB_F32_SHAPES: list[tuple[int, int, str]] = [ + (1, 8, "single 32B block"), + (1, F32_EPR, "single-repeat"), + (1, F32_EPR + 32, "tail repeat padded offsets"), + (2, F32_EPR, "multi-repeat 2 rows"), + (1, F32_EPR * 2, "multi-repeat 2 repeats"), + (4, F32_EPR, "multi-repeat 4 rows"), +] + +GATHERB_F32_PARAMS = [ + pytest.param( + (rows, cols, desc), + id=f"gatherb-float32-{rows}x{cols}-{desc.replace(' ', '-')}", + ) + for rows, cols, desc in GATHERB_F32_SHAPES +] + +GATHERB_F16_SHAPES: list[tuple[int, int, str]] = [ + (1, 16, "single 32B block"), + (1, F16_EPR, "single-repeat"), + (1, F16_EPR + 64, "tail repeat padded offsets"), +] + +GATHERB_F16_PARAMS = [ + pytest.param( + (rows, cols, desc), + id=f"gatherb-float16-{rows}x{cols}-{desc.replace(' ', '-')}", + ) + for rows, cols, desc in GATHERB_F16_SHAPES +] + +GATHER_INDEX_SHAPES: list[tuple[int, int, str]] = [ + (1, 64, "single-repeat"), + (2, 64, "multi-row"), + (2, 128, "multi-row wide"), +] + +GATHER_INDEX_PARAMS = [ + pytest.param( + (rows, cols, desc, dtype_str), + id=f"gather-index-{dtype_str}-{rows}x{cols}-{desc.replace(' ', '-')}", + ) + for dtype_str in ("float32", "float16") + for rows, cols, desc in GATHER_INDEX_SHAPES +] + + +@pytest.mark.require_npu +@pytest.mark.parametrize("case", GATHERB_F32_PARAMS) +def test_gatherb_f32(case, torch, target_arch, backend): + rows, cols, desc = case + + kernel = make_gatherb_kernel( + rows, cols, dtype_str="float32", + target=target_arch, backend=backend, + ) + compile_s, launch_s = launch_and_check_gather( + rows=rows, + cols=cols, + dtype_str="float32", + torch=torch, + kernel_handle=kernel, + ) + print(f" PASS gatherb float32 {rows}x{cols} ({desc}) " + f"compile={compile_s:.3f}s launch={launch_s:.3f}s") + + +@pytest.mark.require_npu +@pytest.mark.parametrize("case", GATHERB_F16_PARAMS) +def test_gatherb_f16(case, torch, target_arch, backend): + rows, cols, desc = case + + kernel = make_gatherb_kernel( + rows, cols, dtype_str="float16", + target=target_arch, backend=backend, + ) + compile_s, launch_s = launch_and_check_gather( + rows=rows, + cols=cols, + dtype_str="float16", + torch=torch, + kernel_handle=kernel, + ) + print(f" PASS gatherb float16 {rows}x{cols} ({desc}) " + f"compile={compile_s:.3f}s launch={launch_s:.3f}s") + + +@pytest.mark.require_npu +@pytest.mark.parametrize("case", GATHER_INDEX_PARAMS) +def test_gather_index(case, torch, target_arch, backend): + rows, cols, desc, dtype_str = case + + kernel = make_gather_index_kernel( + rows, cols, dtype_str=dtype_str, + target=target_arch, backend=backend, + ) + compile_s, launch_s = launch_and_check_gather_index( + rows=rows, + cols=cols, + dtype_str=dtype_str, + torch=torch, + kernel_handle=kernel, + ) + print(f" PASS gather index {dtype_str} {rows}x{cols} ({desc}) " + f"compile={compile_s:.3f}s launch={launch_s:.3f}s") diff --git a/test/lit/pto/tgatherb_a2a3_offsets_dtype_invalid.pto b/test/lit/pto/tgatherb_a2a3_offsets_dtype_invalid.pto new file mode 100644 index 0000000000..bcdf445106 --- /dev/null +++ b/test/lit/pto/tgatherb_a2a3_offsets_dtype_invalid.pto @@ -0,0 +1,20 @@ +// 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 software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-backend=vpto %s 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @bad_offset_dtype() attributes {pto.aicore} { + %src = pto.alloc_tile : !pto.tile_buf + %off = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: expects offsets element type to be i32 + pto.tgatherb ins(%src, %off : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tgatherb_a2a3_offsets_invalid.pto b/test/lit/pto/tgatherb_a2a3_offsets_invalid.pto new file mode 100644 index 0000000000..f4671b6d04 --- /dev/null +++ b/test/lit/pto/tgatherb_a2a3_offsets_invalid.pto @@ -0,0 +1,20 @@ +// 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 software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-backend=vpto %s 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @bad_offset_cols() attributes {pto.aicore} { + %src = pto.alloc_tile : !pto.tile_buf + %off = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: expects offsets valid cols to be compact 32B block address count padded to 8 entries; expected 8, got 64 + pto.tgatherb ins(%src, %off : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/tile_to_ub/gather/tgather_index_a3.pto b/test/lit/vpto/ub/tile_to_ub/gather/tgather_index_a3.pto new file mode 100644 index 0000000000..5eec34044f --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/tgather_index_a3.pto @@ -0,0 +1,27 @@ +// 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 software repository for the full text of the License. + +// Slice 2: pto.tgather (index form) on a3/vpto lowers through the tile→UB +// pipeline via index→byte-offset decomposition (vmuls + vgather with srcAddr +// packed as offsetAddr), matching the pto-isa a2a3/TGather.hpp conversion. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=pto-lower-to-ubuf-ops 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: func.func @lower_tgather_index_a3() + func.func @lower_tgather_index_a3() attributes {pto.aicore} { + %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 + // CHECK: pto.ub.vgather + // CHECK-NOT: pto.tgather + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_a3.pto b/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_a3.pto new file mode 100644 index 0000000000..bf7958cffd --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_a3.pto @@ -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 software repository for the full text of the License. + +// tgatherb lowers to pto.ub.vgatherb on a3/vpto. The offset tile is compact: +// one i32 32B-block address per output 32B block (8 entries for 64 f32 values). + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=pto-lower-to-ubuf-ops 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: func.func @lower_tgatherb_a3() + func.func @lower_tgatherb_a3() attributes {pto.aicore} { + %src = pto.alloc_tile : !pto.tile_buf + %off = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: pto.ub.vgatherb + // CHECK-NOT: pto.tgatherb + pto.tgatherb ins(%src, %off : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_multi_repeat.pto b/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_multi_repeat.pto new file mode 100644 index 0000000000..3510f6bd20 --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/tgatherb_multi_repeat.pto @@ -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 software repository for the full text of the License. + +// Multi-row tgatherb: 2×64 f32 tile. Each row uses 8 compact 32B block +// addresses, not 64 scalar offsets. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=pto-lower-to-ubuf-ops 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: func.func @tgatherb_multi_repeat + func.func @tgatherb_multi_repeat() attributes {pto.aicore} { + %src = pto.alloc_tile : !pto.tile_buf + %off = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: pto.ub.vgatherb + // CHECK-NOT: pto.tgatherb + pto.tgatherb ins(%src, %off : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/ub_to_llvm/gather/tgather_index_full_chain.pto b/test/lit/vpto/ub/ub_to_llvm/gather/tgather_index_full_chain.pto new file mode 100644 index 0000000000..cb8cb927e0 --- /dev/null +++ b/test/lit/vpto/ub/ub_to_llvm/gather/tgather_index_full_chain.pto @@ -0,0 +1,27 @@ +// 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 software repository for the full text of the License. + +// Full-chain: pto.tgather (index form) → vmuls + ub.vgather → VGATHER.b32. +// Verifies the complete MLIR lowering pipeline for the element-index gather. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: llvm.func @tgather_index_full_chain + func.func @tgather_index_full_chain() attributes {pto.aicore} { + %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 + // CHECK: llvm.call @llvm.hivm.VGATHER.b32 + // CHECK-NOT: pto.tgather + // CHECK-NOT: pto.ub. + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/ub_to_llvm/gather/tgatherb_full_chain.pto b/test/lit/vpto/ub/ub_to_llvm/gather/tgatherb_full_chain.pto new file mode 100644 index 0000000000..026b024daa --- /dev/null +++ b/test/lit/vpto/ub/ub_to_llvm/gather/tgatherb_full_chain.pto @@ -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 software repository for the full text of the License. + +// Full-chain: pto.tgatherb → pto.ub.vgatherb → llvm.hivm.VGATHERB.b32. +// Verifies the complete MLIR lowering pipeline composes for a3/vpto. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: llvm.func @tgatherb_full_chain + func.func @tgatherb_full_chain() attributes {pto.aicore} { + %src = pto.alloc_tile : !pto.tile_buf + %off = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: llvm.call @llvm.hivm.VGATHERB.b32({{%[^,]+}}, {{%[^,]+}}, {{%[^)]+}}) + // CHECK-NOT: pto.tgatherb + // CHECK-NOT: pto.ub.vgatherb + pto.tgatherb ins(%src, %off : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/ub/ub_to_llvm/gather/vgather_dtypes.pto b/test/lit/vpto/ub/ub_to_llvm/gather/vgather_dtypes.pto new file mode 100644 index 0000000000..40a0c289ac --- /dev/null +++ b/test/lit/vpto/ub/ub_to_llvm/gather/vgather_dtypes.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the software repository for the full text of the License. + +// pto.ub.vgather -> llvm.hivm.VGATHER.b16/.b32 (packed config form). + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: llvm.func @ub_vgather_f32 + func.func @ub_vgather_f32(%dst: !pto.ptr, %src: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : i64 + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + // CHECK: llvm.call @llvm.hivm.VGATHER.b32({{%[^,]+}}, {{%[^,]+}}, {{%[^)]+}}) : (!llvm.ptr<6>, !llvm.ptr<6>, i64) -> () + // CHECK-NOT: pto.ub. + pto.ub.vgather %dst, %src, %c0, %c8, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64 + return + } + + // CHECK-LABEL: llvm.func @ub_vgather_f16 + func.func @ub_vgather_f16(%dst: !pto.ptr, %src: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : i64 + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + // CHECK: llvm.call @llvm.hivm.VGATHER.b16({{%[^,]+}}, {{%[^,]+}}, {{%[^)]+}}) : (!llvm.ptr<6>, !llvm.ptr<6>, i64) -> () + // CHECK-NOT: pto.ub. + pto.ub.vgather %dst, %src, %c0, %c8, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64 + return + } +} diff --git a/test/lit/vpto/ub/ub_to_llvm/gather/vgatherb_dtypes.pto b/test/lit/vpto/ub/ub_to_llvm/gather/vgatherb_dtypes.pto new file mode 100644 index 0000000000..0421264358 --- /dev/null +++ b/test/lit/vpto/ub/ub_to_llvm/gather/vgatherb_dtypes.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the software repository for the full text of the License. + +// pto.ub.vgatherb -> llvm.hivm.VGATHERB.b16/.b32 (packed config form). +// Verifies: element-width suffix (b32/b16), 3-operand packed form +// (dst_ptr, offset_ptr, i64 config), and that the op is fully lowered. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + // CHECK-LABEL: llvm.func @ub_vgatherb_f32 + func.func @ub_vgatherb_f32(%dst: !pto.ptr, %off: !pto.ptr, %src: !pto.ptr) attributes {pto.aicore} { + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + // CHECK: llvm.call @llvm.hivm.VGATHERB.b32({{%[^,]+}}, {{%[^,]+}}, {{%[^)]+}}) : (!llvm.ptr<6>, !llvm.ptr<6>, i64) -> () + // CHECK-NOT: pto.ub. + pto.ub.vgatherb %dst, %off, %src, %c8, %c1, %c1 : !pto.ptr, !pto.ptr, !pto.ptr, i64, i64, i64 + return + } + + // CHECK-LABEL: llvm.func @ub_vgatherb_f16 + func.func @ub_vgatherb_f16(%dst: !pto.ptr, %off: !pto.ptr, %src: !pto.ptr) attributes {pto.aicore} { + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + // CHECK: llvm.call @llvm.hivm.VGATHERB.b16({{%[^,]+}}, {{%[^,]+}}, {{%[^)]+}}) : (!llvm.ptr<6>, !llvm.ptr<6>, i64) -> () + // CHECK-NOT: pto.ub. + pto.ub.vgatherb %dst, %off, %src, %c8, %c1, %c1 : !pto.ptr, !pto.ptr, !pto.ptr, i64, i64, i64 + return + } +} diff --git a/test/lit/vpto/ub/ub_vgather_parse.pto b/test/lit/vpto/ub/ub_vgather_parse.pto new file mode 100644 index 0000000000..5f1d75f3ec --- /dev/null +++ b/test/lit/vpto/ub/ub_vgather_parse.pto @@ -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 software repository for the full text of the License. + +// Round-trip / parse-verify for the A2/A3 UB element gather primitive pto.ub.vgather. + +// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s -o - 2>&1 | FileCheck %s + +module { + // CHECK-LABEL: func.func @ub_vgather_parse + func.func @ub_vgather_parse() { + %c0 = arith.constant 0 : i64 + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + %dst = pto.castptr %c0 : i64 -> !pto.ptr + %src = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: pto.ub.vgather + // CHECK: !pto.ptr, !pto.ptr, i64, i64, i64 + pto.ub.vgather %dst, %src, %c0, %c8, %c1 : !pto.ptr, !pto.ptr, i64, i64, i64 + return + } +} diff --git a/test/lit/vpto/ub/ub_vgatherb_parse.pto b/test/lit/vpto/ub/ub_vgatherb_parse.pto new file mode 100644 index 0000000000..b55698f097 --- /dev/null +++ b/test/lit/vpto/ub/ub_vgatherb_parse.pto @@ -0,0 +1,29 @@ +// 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 software repository for the full text of the License. + +// Round-trip / parse-verify for the A2/A3 UB gather primitive pto.ub.vgatherb. +// Lowers to the packed CCE vgatherb intrinsic (llvm.hivm.VGATHERB.b16/.b32); +// distinct from the A5 vreg pto.vgather2 path. + +// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s -o - 2>&1 | FileCheck %s + +module { + // CHECK-LABEL: func.func @ub_vgatherb_parse + func.func @ub_vgatherb_parse() { + %c0 = arith.constant 0 : i64 + %c8 = arith.constant 8 : i64 + %c1 = arith.constant 1 : i64 + %dst = pto.castptr %c0 : i64 -> !pto.ptr + %off = pto.castptr %c0 : i64 -> !pto.ptr + %src = pto.castptr %c0 : i64 -> !pto.ptr + // CHECK: pto.ub.vgatherb + // CHECK: !pto.ptr, !pto.ptr, !pto.ptr, i64, i64, i64 + pto.ub.vgatherb %dst, %off, %src, %c8, %c1, %c1 : !pto.ptr, !pto.ptr, !pto.ptr, i64, i64, i64 + return + } +}