From 62447c1a071f80146cf0e54f111316c78c59835e Mon Sep 17 00:00:00 2001 From: Giacomo Castiglioni Date: Thu, 23 Jul 2026 18:21:55 +0200 Subject: [PATCH 1/2] feat(vpto): add A2/A3 mgather lowering --- docs/designs/a2a3-vpto-tgather.md | 39 ++- lib/PTO/Transforms/LowerPTOToUBufOps.cpp | 331 ++++++++++++++++-- .../Transforms/PTOMaterializeTileHandles.cpp | 7 + lib/PTO/Transforms/VPTOPtrCastCleanup.cpp | 15 +- .../ub/tile_to_ub/gather/mgather_a2a3.pto | 87 +++++ .../gather/mgather_a2a3_gm_l1_unsupported.pto | 23 ++ .../tile_to_ub/gather/mgather_a2a3_types.pto | 72 ++++ tools/ptoas/ptoas.cpp | 1 + 8 files changed, 538 insertions(+), 37 deletions(-) create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_gm_l1_unsupported.pto create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_types.pto diff --git a/docs/designs/a2a3-vpto-tgather.md b/docs/designs/a2a3-vpto-tgather.md index 1a07d9273c..0589573ee3 100644 --- a/docs/designs/a2a3-vpto-tgather.md +++ b/docs/designs/a2a3-vpto-tgather.md @@ -1,7 +1,8 @@ # 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, +Scope: bring tile-level gather (`pto.tgather`, `pto.tgatherb`, and the GM-to-UB +form of `pto.mgather`) 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. @@ -65,10 +66,10 @@ identical to `LowerPTOToUBufOps`'s `dispatch` head/tail repeat loop. ### Related A2/A3 operations - **`MGather` exists on A2/A3** — `npu/a2a3/MGather.hpp` implements GM gather - loads for row and element coalescing. Its GM-to-UB VPTO lowering is handled - as a separate follow-up because it decomposes into scalar address generation, - scalar loads/stores, and per-row MTE2 copies rather than the UB-local - `vgather`/`vgatherb` instructions covered here. + loads for row and element coalescing. Its GM-to-UB VPTO lowering is implemented + as a stacked follow-up using scalar address generation, scalar loads/stores, + and per-row MTE2 copies rather than the UB-local `vgather`/`vgatherb` + instructions covered here. - **No vreg** — A5's `vgather2`/`vgather2_bc` (`npu/a5/TGather.hpp:51,79,108`) must not be used on A2/A3. @@ -82,11 +83,28 @@ lowering must reproduce: |---|---|---| | `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`) | Supported by A2/A3 PTO-ISA/EmitC; VPTO lowering is a separate follow-up | +| `pto.mgather` | `emitc::CallOpaqueOp "MGATHER"` (`:3323`) | GM-to-UB Row/Elem lowering is implemented for A2/A3 VPTO; GM-to-L1 remains deferred | MGATHER semantics: `dst[r,:]=table[idx[r],:]` (Coalesce::Row) or `dst[i,j]=table[idx[i,j]]` (Coalesce::Elem), with `GatherOOB` modes. +### A2/A3 VPTO MGATHER decomposition + +`LowerPTOToUBufOps` lowers GM-to-UB MGATHER without introducing a new UB +intrinsic: + +- `Coalesce::Row` loads each `i32` index on the scalar pipe, normalizes it for + `Undefined`, `Clamp`, `Wrap`, or `Zero`, then issues a row-sized + `pto.mte_gm_ub`. The scalar-to-MTE2 handoff uses the macro-reserved + `EVENT_ID0`. +- `Coalesce::Elem` performs scalar indexed GM loads and UB stores for each valid + destination element, with the same four OOB policies. +- The lowering accepts the A2/A3 payload set + (`i8/ui8/i16/ui16/i32/ui32/f16/bf16/f32`) and preserves arbitrary row-major + GM view strides by flattening logical indices through extracted metadata. +- GM-to-L1 MGATHER is rejected explicitly at the A2/A3 VPTO lowering boundary + until its separate NZ/scratch path is implemented. + ## Current `pto.ub.*` layer Defined in `include/PTO/IR/VPTOUbOps.td`. Before this work it was @@ -219,9 +237,10 @@ job for these today. 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. **Separate follow-up**: A2/A3 `mgather` GM-to-UB Row/Elem lowering. This - requires scalar index/OOB handling plus MTE2 or scalar GM access and is kept - separate from the UB-local gather implementation in this document. +4. **Stacked follow-up — DONE**: A2/A3 `mgather` GM-to-UB Row/Elem lowering, + including scalar index handling, all four `GatherOOB` modes, per-row MTE2 + copies, scalar element access, the full A2/A3 payload set, and a fail-closed + GM-to-L1 diagnostic. ## References diff --git a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp index f2d66ae5bd..f0e31a7176 100644 --- a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp +++ b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp @@ -65,6 +65,13 @@ static unsigned getElementSize(Type elemTy) { return 0; } +static unsigned getMGatherElementSize(Type elemTy) { + if (auto intTy = dyn_cast(elemTy); + intTy && intTy.getWidth() == 8) + return 1; + return getElementSize(elemTy); +} + static Type getStoredElemType(Type ty) { if (auto tbTy = dyn_cast(ty)) return tbTy.getElementType(); @@ -138,6 +145,8 @@ struct TileShapeInfo { struct TileShapeMetadata { SmallVector shape; SmallVector validShape; + pto::AddressSpace memorySpace; + pto::BLayout bLayout; }; using TileShapeMap = DenseMap; @@ -266,7 +275,11 @@ struct LowerPTOToUBufOpsPass auto pc = builder.create(op.getLoc(), ptrTy, addr); tileShapes[pc.getResult()] = { SmallVector(shape), - SmallVector(tbTy.getValidShape())}; + SmallVector(tbTy.getValidShape()), + cast(tbTy.getMemorySpace()) + .getAddressSpace(), + tbTy.getConfigAttr() ? tbTy.getConfigAttr().getBLayout().getValue() + : pto::BLayout::RowMajor}; op.getResult().replaceAllUsesWith(pc.getResult()); op.erase(); } @@ -1009,6 +1022,20 @@ struct LowerPTOToUBufOpsPass } } + // ---- mgather (GM -> UB) ---- + { + SmallVector ops; + func.walk([&](pto::MGatherOp op) { ops.push_back(op); }); + for (auto op : ops) { + builder.setInsertionPoint(op); + if (failed(lowerMGather(op, builder, tileShapes))) { + signalPassFailure(); + return; + } + op.erase(); + } + } + // ---- cleanup dead PTO ops ---- SmallVector toErase; func.walk([&](Operation *op) { @@ -1438,8 +1465,8 @@ struct LowerPTOToUBufOpsPass return !strides.empty() && strides.back() == 1; } - static FailureOr extractDmaViewInfo(pto::TLoadOp op) { - auto pvOp = op.getSrc().getDefiningOp(); + static FailureOr extractDmaViewInfo(Operation *op, Value view) { + auto pvOp = view.getDefiningOp(); if (pvOp) { auto mtvOp = pvOp.getSource().getDefiningOp(); if (!mtvOp) @@ -1451,34 +1478,21 @@ struct LowerPTOToUBufOpsPass mtvOp.getStrides().end()); if (info.strides.empty() || !matchPattern(info.strides.back(), m_One())) { - op.emitError("A2/A3 DMA lowering requires a unit innermost stride"); + op->emitError("A2/A3 DMA lowering requires a unit innermost stride"); return failure(); } info.offsets.assign(pvOp.getOffsets().begin(), pvOp.getOffsets().end()); return info; } - return extractDmaMemRefViewInfo(op.getLoc(), op.getSrc(), op.getContext()); + return extractDmaMemRefViewInfo(op->getLoc(), view, op->getContext()); + } + + static FailureOr extractDmaViewInfo(pto::TLoadOp op) { + return extractDmaViewInfo(op.getOperation(), op.getSrc()); } static FailureOr extractDmaViewInfo(pto::TStoreOp op) { - auto pvOp = op.getDst().getDefiningOp(); - if (pvOp) { - auto mtvOp = pvOp.getSource().getDefiningOp(); - if (!mtvOp) - return failure(); - DmaViewInfo info; - info.gmPtr = mtvOp.getPtr(); - info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); - info.strides.assign(mtvOp.getStrides().begin(), mtvOp.getStrides().end()); - if (info.strides.empty() || - !matchPattern(info.strides.back(), m_One())) { - op.emitError("A2/A3 DMA lowering requires a unit innermost stride"); - return failure(); - } - info.offsets.assign(pvOp.getOffsets().begin(), pvOp.getOffsets().end()); - return info; - } - return extractDmaMemRefViewInfo(op.getLoc(), op.getDst(), op.getContext()); + return extractDmaViewInfo(op.getOperation(), op.getDst()); } static FailureOr extractDmaMemRefViewInfo(Location loc, Value view, @@ -1490,7 +1504,7 @@ struct LowerPTOToUBufOpsPass if (!msAttr || msAttr.getAddressSpace() != pto::AddressSpace::GM) return failure(); ArrayRef shape = memTy.getShape(); - if (shape.size() < 2) + if (shape.empty()) return failure(); if (!hasUnitInnermostStride(view)) { emitError(loc) << "A2/A3 DMA lowering requires a unit innermost stride"; @@ -1503,7 +1517,8 @@ struct LowerPTOToUBufOpsPass auto ptrTy = pto::PtrType::get(ctx, memTy.getElementType(), msAttr); auto metadata = b.create(loc, view); info.gmPtr = b.create(loc, ptrTy, traceRootMemRef(view)); - info.linearOffset = metadata.getOffset(); + if (memTy.getStridesAndOffset().second != 0) + info.linearOffset = metadata.getOffset(); info.sizes.assign(metadata.getSizes().begin(), metadata.getSizes().end()); info.strides.assign(metadata.getStrides().begin(), metadata.getStrides().end()); @@ -1562,13 +1577,17 @@ struct LowerPTOToUBufOpsPass auto origPtrTy = cast(gmPtr.getType()); auto bytePtrTy = pto::PtrType::get(b.getContext(), b.getI8Type(), origPtrTy.getMemorySpace()); - Value bytePtr = b.create(loc, bytePtrTy, gmPtr); + Value bytePtr = gmPtr; + if (bytePtrTy != origPtrTy) + bytePtr = b.create(loc, bytePtrTy, gmPtr); Value offIdx = byteOff; if (!offIdx.getType().isIndex()) offIdx = b.create(loc, b.getIndexType(), byteOff) .getResult(); Value offsetBytePtr = b.create(loc, bytePtrTy, bytePtr, offIdx); + if (bytePtrTy == origPtrTy) + return offsetBytePtr; return b.create(loc, origPtrTy, offsetBytePtr); } @@ -1669,6 +1688,266 @@ struct LowerPTOToUBufOpsPass return success(); } + struct NormalizedGatherIndex { + Value value; + Value inBounds; + }; + + Value product(Location loc, OpBuilder &b, ArrayRef values) { + Value result = idxc1(loc, b); + for (Value value : values) + result = b.create(loc, result, value); + return result; + } + + NormalizedGatherIndex normalizeGatherIndex(Location loc, OpBuilder &b, + Value rawIndex, Value count, + pto::GatherOOB mode) { + Value index = b.create(loc, b.getIndexType(), rawIndex); + Value zero = idxc0(loc, b); + Value negative = + b.create(loc, arith::CmpIPredicate::slt, index, zero); + Value tooLarge = + b.create(loc, arith::CmpIPredicate::sge, index, count); + Value outOfBounds = b.create(loc, negative, tooLarge); + Value inBounds = b.create( + loc, outOfBounds, b.create(loc, 1, 1)); + + switch (mode) { + case pto::GatherOOB::Undefined: + case pto::GatherOOB::Zero: + return {index, inBounds}; + case pto::GatherOOB::Clamp: { + Value nonNegative = b.create(loc, negative, zero, index); + Value last = b.create(loc, count, idxc1(loc, b)); + Value aboveLast = b.create(loc, arith::CmpIPredicate::sgt, + nonNegative, last); + return {b.create(loc, aboveLast, last, nonNegative), + inBounds}; + } + case pto::GatherOOB::Wrap: { + Value remainder = b.create(loc, index, count); + Value wrapped = b.create(loc, remainder, count); + Value remainderNegative = b.create( + loc, arith::CmpIPredicate::slt, remainder, zero); + return { + b.create(loc, remainderNegative, wrapped, remainder), + inBounds}; + } + } + llvm_unreachable("unknown MGATHER OOB mode"); + } + + Value flattenGatherIndex(Location loc, OpBuilder &b, Value flatIndex, + ArrayRef sizes, ArrayRef strides) { + Value remaining = flatIndex; + Value linearOffset = idxc0(loc, b); + for (int64_t i = static_cast(sizes.size()) - 1; i >= 0; --i) { + Value coordinate = b.create(loc, remaining, sizes[i]); + Value contribution = b.create(loc, coordinate, strides[i]); + linearOffset = b.create(loc, linearOffset, contribution); + if (i != 0) + remaining = b.create(loc, remaining, sizes[i]); + } + return linearOffset; + } + + void emitScalarZero(Location loc, OpBuilder &b, Value dst, Value dstOffset, + Type elemTy) { + Value zero = + b.create(loc, elemTy, b.getZeroAttr(elemTy)); + b.create(loc, dst, dstOffset, zero); + } + + void emitScalarGather(Location loc, OpBuilder &b, Value gmPtr, Value dst, + Value dstOffset, Value sourceOffset, Type elemTy) { + Value value = b.create(loc, elemTy, gmPtr, sourceOffset); + b.create(loc, dst, dstOffset, value); + } + + void emitScalarToMte2Sync(Location loc, OpBuilder &b) { + auto pipeS = pto::PipeAttr::get(b.getContext(), pto::PIPE::PIPE_S); + auto pipeMte2 = pto::PipeAttr::get(b.getContext(), pto::PIPE::PIPE_MTE2); + auto event0 = + pto::EventAttr::get(b.getContext(), static_cast(0)); + b.create(loc, pipeS, pipeMte2, event0); + b.create(loc, pipeS, pipeMte2, event0); + } + + LogicalResult emitMGatherRow(pto::MGatherOp op, OpBuilder &b, + const DmaViewInfo &viewInfo, Value gmPtr, + Value idxPtr, Value dstPtr, + const TileShapeMetadata &idxMeta, + const TileShapeMetadata &dstMeta, Type elemTy, + pto::GatherOOB oob) { + Location loc = op.getLoc(); + if (viewInfo.sizes.size() < 2) + return op.emitOpError( + "requires a rank-2 or greater GM table for coalesce=row"); + + ArrayRef rowSizes(viewInfo.sizes.data(), viewInfo.sizes.size() - 1); + ArrayRef rowStrides(viewInfo.strides.data(), + viewInfo.strides.size() - 1); + Value tableRows = product(loc, b, rowSizes); + int64_t dstRows = dstMeta.validShape[0]; + int64_t dstCols = dstMeta.validShape[1]; + int64_t dstStride = dstMeta.shape[1]; + bool contiguousIndex = + idxMeta.validShape[0] <= 1 || idxMeta.bLayout == pto::BLayout::ColMajor; + unsigned elemSize = getMGatherElementSize(elemTy); + + for (int64_t row = 0; row < dstRows; ++row) { + Value idxOffset = + idxc(contiguousIndex ? row : row * idxMeta.shape[1], loc, b); + Value rawIndex = + b.create(loc, b.getI32Type(), idxPtr, idxOffset); + NormalizedGatherIndex index = + normalizeGatherIndex(loc, b, rawIndex, tableRows, oob); + Value dstOffset = idxc(row * dstStride, loc, b); + + auto emitCopy = [&]() { + Value sourceOffset = + flattenGatherIndex(loc, b, index.value, rowSizes, rowStrides); + Value source = + b.create(loc, gmPtr.getType(), gmPtr, sourceOffset); + Value destination = + b.create(loc, dstPtr.getType(), dstPtr, dstOffset); + emitScalarToMte2Sync(loc, b); + pto::DmaLoopConfig nburst{i64c1(loc, b), i64c0(loc, b), i64c0(loc, b)}; + b.create(loc, source, destination, i64c0(loc, b), + i64c(dstCols * elemSize, loc, b), nburst, + llvm::ArrayRef{}, + std::nullopt); + }; + + if (oob != pto::GatherOOB::Zero) { + emitCopy(); + continue; + } + + auto ifOp = b.create(loc, TypeRange{}, index.inBounds, + /*addThenBlock=*/true, + /*addElseBlock=*/true); + b.setInsertionPointToStart(ifOp.thenBlock()); + emitCopy(); + b.create(loc); + b.setInsertionPointToStart(ifOp.elseBlock()); + for (int64_t col = 0; col < dstCols; ++col) + emitScalarZero(loc, b, dstPtr, idxc(row * dstStride + col, loc, b), + elemTy); + emitScalarToMte2Sync(loc, b); + b.create(loc); + b.setInsertionPointAfter(ifOp); + } + return success(); + } + + LogicalResult emitMGatherElem(pto::MGatherOp op, OpBuilder &b, + const DmaViewInfo &viewInfo, Value gmPtr, + Value idxPtr, Value dstPtr, + const TileShapeMetadata &idxMeta, + const TileShapeMetadata &dstMeta, Type elemTy, + pto::GatherOOB oob) { + Location loc = op.getLoc(); + if (viewInfo.sizes.empty()) + return op.emitOpError("requires a non-empty GM table"); + + Value tableElements = product(loc, b, viewInfo.sizes); + int64_t dstRows = dstMeta.validShape[0]; + int64_t dstCols = dstMeta.validShape[1]; + int64_t dstStride = dstMeta.shape[1]; + int64_t idxStride = idxMeta.shape[1]; + for (int64_t row = 0; row < dstRows; ++row) { + for (int64_t col = 0; col < dstCols; ++col) { + Value idxOffset = idxc(row * idxStride + col, loc, b); + Value dstOffset = idxc(row * dstStride + col, loc, b); + Value rawIndex = + b.create(loc, b.getI32Type(), idxPtr, idxOffset); + NormalizedGatherIndex index = + normalizeGatherIndex(loc, b, rawIndex, tableElements, oob); + Value sourceOffset = flattenGatherIndex( + loc, b, index.value, viewInfo.sizes, viewInfo.strides); + + if (oob != pto::GatherOOB::Zero) { + emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, sourceOffset, + elemTy); + continue; + } + + auto ifOp = b.create(loc, TypeRange{}, index.inBounds, + /*addThenBlock=*/true, + /*addElseBlock=*/true); + b.setInsertionPointToStart(ifOp.thenBlock()); + emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, sourceOffset, + elemTy); + b.create(loc); + b.setInsertionPointToStart(ifOp.elseBlock()); + emitScalarZero(loc, b, dstPtr, dstOffset, elemTy); + b.create(loc); + b.setInsertionPointAfter(ifOp); + } + } + return success(); + } + + LogicalResult lowerMGather(pto::MGatherOp op, OpBuilder &b, + const TileShapeMap &tileShapes) { + auto dstIt = tileShapes.find(op.getDst()); + auto idxIt = tileShapes.find(op.getIdx()); + if (dstIt == tileShapes.end()) + return op.emitOpError( + "A2/A3 VPTO lowering requires an alloc_tile-backed mgather dst"); + if (dstIt->second.memorySpace != pto::AddressSpace::VEC) + return op.emitOpError( + "A2/A3 VPTO lowering currently supports only GM->UB mgather; " + "GM->L1 is deferred"); + if (idxIt == tileShapes.end()) + return op.emitOpError( + "A2/A3 GM->UB VPTO mgather requires an alloc_tile-backed idx"); + if (op.getScratch()) + return op.emitOpError( + "A2/A3 GM->UB VPTO mgather does not support a scratch operand"); + if (dstIt->second.shape.size() != 2 || + dstIt->second.validShape.size() != 2 || + idxIt->second.shape.size() != 2 || idxIt->second.validShape.size() != 2) + return op.emitOpError("requires rank-2 idx and dst tiles"); + + Type elemTy = getStoredElemType(op.getDst().getType()); + if (!elemTy || getMGatherElementSize(elemTy) == 0) + return op.emitOpError("has an unsupported destination element type"); + auto idxPtrTy = dyn_cast(op.getIdx().getType()); + if (!idxPtrTy || !idxPtrTy.getElementType().isInteger(32)) + return op.emitOpError("requires an i32 UB index tile"); + + auto viewInfo = extractDmaViewInfo(op.getOperation(), op.getMem()); + if (failed(viewInfo)) + return op.emitOpError("requires a supported GM tensor view"); + if (viewInfo->sizes.size() != viewInfo->strides.size()) + return op.emitOpError("requires one stride per GM table dimension"); + + Value byteOff = computeGMByteOffset(op.getLoc(), b, *viewInfo, + getMGatherElementSize(elemTy)); + Value gmPtr = offsetGMPtrByBytes(op.getLoc(), b, viewInfo->gmPtr, byteOff); + + auto coalesceAttr = + dyn_cast_or_null(op.getProperties().coalesce); + pto::Coalesce coalesce; + if (coalesceAttr) { + coalesce = coalesceAttr.getValue(); + } else if (idxIt->second.validShape == dstIt->second.validShape) { + coalesce = pto::Coalesce::Elem; + } else { + coalesce = pto::Coalesce::Row; + } + + pto::GatherOOB oob = op.getGatherOob(); + if (coalesce == pto::Coalesce::Row) + return emitMGatherRow(op, b, *viewInfo, gmPtr, op.getIdx(), op.getDst(), + idxIt->second, dstIt->second, elemTy, oob); + return emitMGatherElem(op, b, *viewInfo, gmPtr, op.getIdx(), op.getDst(), + idxIt->second, dstIt->second, elemTy, oob); + } + LogicalResult lowerTLoad(pto::TLoadOp op, OpBuilder &b, const TileShapeMap &tileShapes) { Location loc = op.getLoc(); diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 80dc8ffc51..c382780cca 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -97,6 +97,13 @@ static bool shouldMaterializeOperand(Operation *owner) { if (isa(owner)) return false; + if (isa(owner)) { + auto module = owner->getParentOfType(); + auto backend = module ? module->getAttrOfType("pto.backend") + : StringAttr{}; + return backend && backend.getValue() == "vpto"; + } + StringRef name = owner->getName().getStringRef(); if (name == "pto.set_validshape") return true; diff --git a/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp b/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp index fc238fac5f..7a5721374e 100644 --- a/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp +++ b/lib/PTO/Transforms/VPTOPtrCastCleanup.cpp @@ -25,6 +25,18 @@ using namespace mlir; namespace { +struct EraseTrivialCastPtrPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(pto::CastPtrOp op, + PatternRewriter &rewriter) const override { + if (op.getInput().getType() != op.getResult().getType()) + return failure(); + rewriter.replaceOp(op, op.getInput()); + return success(); + } +}; + struct CollapsePtrMemRefPtrBridgePattern : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -68,7 +80,8 @@ struct VPTOPtrCastCleanupPass void runOnOperation() override { RewritePatternSet patterns(&getContext()); - patterns.add(&getContext()); + patterns.add( + &getContext()); if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) signalPassFailure(); } diff --git a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto new file mode 100644 index 0000000000..d32b934288 --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto @@ -0,0 +1,87 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto-llvm-ir %s -o %t --mlir-print-ir-after=pto-lower-to-ubuf-ops 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a2 --pto-backend=vpto --emit-vpto-llvm-ir %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 @mgather_a2a3 + func.func @mgather_a2a3(%rowSrc: !pto.ptr, %zeroSrc: !pto.ptr, %clampSrc: !pto.ptr, %wrapSrc: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + + %rowView = pto.make_tensor_view %rowSrc, shape = [%c4, %c16], strides = [%c16, %c1] + : !pto.tensor_view + %rowPart = pto.partition_view %rowView, offsets = [%c0, %c0], sizes = [%c4, %c16] + : !pto.tensor_view -> !pto.partition_tensor_view<4x16xf16> + %rowIdx = pto.alloc_tile + : !pto.tile_buf + %rowDst = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%rowPart, %rowIdx : !pto.partition_tensor_view<4x16xf16>, !pto.tile_buf) + outs(%rowDst : !pto.tile_buf) + {coalesce = #pto} + // CHECK: pto.load_scalar + // CHECK: pto.set_flag[, , ] + // CHECK: pto.wait_flag[, , ] + // CHECK: pto.mte_gm_ub + + %zeroView = pto.make_tensor_view %zeroSrc, shape = [%c4, %c32], strides = [%c32, %c1] + : !pto.tensor_view + %zeroPart = pto.partition_view %zeroView, offsets = [%c0, %c0], sizes = [%c4, %c32] + : !pto.tensor_view -> !pto.partition_tensor_view<4x32xi8> + %zeroIdx = pto.alloc_tile + : !pto.tile_buf + %zeroDst = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%zeroPart, %zeroIdx : !pto.partition_tensor_view<4x32xi8>, !pto.tile_buf) + outs(%zeroDst : !pto.tile_buf) + {coalesce = #pto, gatherOob = #pto} + // CHECK: scf.if + // CHECK: pto.mte_gm_ub + // CHECK: arith.constant 0 : i8 + // CHECK: pto.store_scalar + + %clampView = pto.make_tensor_view %clampSrc, shape = [%c8], strides = [%c1] + : !pto.tensor_view + %clampPart = pto.partition_view %clampView, offsets = [%c0], sizes = [%c8] + : !pto.tensor_view -> !pto.partition_tensor_view<8xf32> + %clampIdx = pto.alloc_tile + : !pto.tile_buf + %clampDst = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%clampPart, %clampIdx : !pto.partition_tensor_view<8xf32>, !pto.tile_buf) + outs(%clampDst : !pto.tile_buf) + {coalesce = #pto, gatherOob = #pto} + // CHECK: arith.select + // CHECK: pto.load_scalar {{.*}} : !pto.ptr -> f32 + // CHECK: pto.store_scalar + + %wrapView = pto.make_tensor_view %wrapSrc, shape = [%c8], strides = [%c1] + : !pto.tensor_view + %wrapPart = pto.partition_view %wrapView, offsets = [%c0], sizes = [%c8] + : !pto.tensor_view -> !pto.partition_tensor_view<8xi32> + %wrapIdx = pto.alloc_tile + : !pto.tile_buf + %wrapDst = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%wrapPart, %wrapIdx : !pto.partition_tensor_view<8xi32>, !pto.tile_buf) + outs(%wrapDst : !pto.tile_buf) + {coalesce = #pto, gatherOob = #pto} + // CHECK: arith.remsi + // CHECK: pto.load_scalar {{.*}} : !pto.ptr -> i32 + // CHECK: pto.store_scalar + // CHECK-NOT: pto.mgather + return + } +} diff --git a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_gm_l1_unsupported.pto b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_gm_l1_unsupported.pto new file mode 100644 index 0000000000..d822e42471 --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_gm_l1_unsupported.pto @@ -0,0 +1,23 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto-llvm-ir %s -o /dev/null 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @mgather_gm_l1_unsupported( + %mem: memref<1x1x1x64x32xf16, #pto.address_space>, + %idx: memref<1x1x1x1x32xi32, #pto.address_space>) attributes {pto.aicore} { + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: error: 'pto.mgather' op A2/A3 VPTO lowering currently supports only GM->UB mgather; GM->L1 is deferred + pto.mgather ins(%mem, %idx : memref<1x1x1x64x32xf16, #pto.address_space>, memref<1x1x1x1x32xi32, #pto.address_space>) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } +} diff --git a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_types.pto b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_types.pto new file mode 100644 index 0000000000..6c8603331e --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_types.pto @@ -0,0 +1,72 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto-llvm-ir %s -o /dev/null +// RUN: ptoas --pto-arch=a2 --pto-backend=vpto --emit-vpto-llvm-ir %s -o /dev/null + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @mgather_a2a3_types( + %s_i8: memref<8xi8, #pto.address_space>, + %s_ui8: memref<8xui8, #pto.address_space>, + %s_i16: memref<8xi16, #pto.address_space>, + %s_ui16: memref<8xui16, #pto.address_space>, + %s_i32: memref<8xi32, #pto.address_space>, + %s_ui32: memref<8xui32, #pto.address_space>, + %s_f16: memref<8xf16, #pto.address_space>, + %s_bf16: memref<8xbf16, #pto.address_space>, + %s_f32: memref<8xf32, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + + %d_i8 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_i8, %idx : memref<8xi8, #pto.address_space>, !pto.tile_buf) + outs(%d_i8 : !pto.tile_buf) {coalesce = #pto} + + %d_ui8 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_ui8, %idx : memref<8xui8, #pto.address_space>, !pto.tile_buf) + outs(%d_ui8 : !pto.tile_buf) {coalesce = #pto} + + %d_i16 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_i16, %idx : memref<8xi16, #pto.address_space>, !pto.tile_buf) + outs(%d_i16 : !pto.tile_buf) {coalesce = #pto} + + %d_ui16 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_ui16, %idx : memref<8xui16, #pto.address_space>, !pto.tile_buf) + outs(%d_ui16 : !pto.tile_buf) {coalesce = #pto} + + %d_i32 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_i32, %idx : memref<8xi32, #pto.address_space>, !pto.tile_buf) + outs(%d_i32 : !pto.tile_buf) {coalesce = #pto} + + %d_ui32 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_ui32, %idx : memref<8xui32, #pto.address_space>, !pto.tile_buf) + outs(%d_ui32 : !pto.tile_buf) {coalesce = #pto} + + %d_f16 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_f16, %idx : memref<8xf16, #pto.address_space>, !pto.tile_buf) + outs(%d_f16 : !pto.tile_buf) {coalesce = #pto} + + %d_bf16 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_bf16, %idx : memref<8xbf16, #pto.address_space>, !pto.tile_buf) + outs(%d_bf16 : !pto.tile_buf) {coalesce = #pto} + + %d_f32 = pto.alloc_tile + : !pto.tile_buf + pto.mgather ins(%s_f32, %idx : memref<8xf32, #pto.address_space>, !pto.tile_buf) + outs(%d_f32 : !pto.tile_buf) {coalesce = #pto} + return + } +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index f31c065829..959d0e8c88 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -2736,6 +2736,7 @@ static void prepareVPTOForEmission(PassManager &pm) { pto::createPTONarrowVPTOLoopCountersPass()); kernelModulePM.addPass(createCanonicalizerPass()); kernelModulePM.addPass(createCSEPass()); + kernelModulePM.addPass(pto::createVPTOPtrCastCleanupPass()); kernelModulePM.addPass(pto::createPTOValidateVPTOEmissionIRPass()); } From a106c4b0265c36bdf12fc5b5efa0ca42a1c98fa4 Mon Sep 17 00:00:00 2001 From: Giacomo Castiglioni Date: Tue, 28 Jul 2026 17:02:03 +0200 Subject: [PATCH 2/2] fix(vpto): harden A2/A3 mgather lowering --- docs/designs/a2a3-vpto-tgather.md | 14 +- lib/PTO/Transforms/LowerPTOToUBufOps.cpp | 254 ++++++++++++++---- ptodsl/ptodsl/_ops.py | 26 ++ ptodsl/ptodsl/_tile_namespace.py | 1 + ptodsl/tests/e2e/common.py | 155 +++++++++++ ptodsl/tests/e2e/test_gather.py | 59 ++++ ptodsl/tests/test_vector_cube_ops.py | 29 ++ .../ub/tile_to_ub/gather/mgather_a2a3.pto | 18 +- .../gather/mgather_a2a3_unsupported.pto | 83 ++++++ 9 files changed, 581 insertions(+), 58 deletions(-) create mode 100644 test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_unsupported.pto diff --git a/docs/designs/a2a3-vpto-tgather.md b/docs/designs/a2a3-vpto-tgather.md index 0589573ee3..18d51bce70 100644 --- a/docs/designs/a2a3-vpto-tgather.md +++ b/docs/designs/a2a3-vpto-tgather.md @@ -100,8 +100,14 @@ intrinsic: - `Coalesce::Elem` performs scalar indexed GM loads and UB stores for each valid destination element, with the same four OOB policies. - The lowering accepts the A2/A3 payload set - (`i8/ui8/i16/ui16/i32/ui32/f16/bf16/f32`) and preserves arbitrary row-major - GM view strides by flattening logical indices through extracted metadata. + (`i8/ui8/i16/ui16/i32/ui32/f16/bf16/f32`) with signless `i32` indices and + row-major, none-box UB tiles. Source and destination element types must match. +- `Coalesce::Row` requires a statically positive effective rank-2 ND table, + exact source/destination row-width agreement, unit innermost stride, and a + 32-byte-aligned row transfer. `Coalesce::Elem` requires a statically positive, + dense contiguous table and treats each index as a direct flat element offset. +- Unsupported dynamic, empty, strided Elem, column-major, row-plus-one, scratch, + and implicit-coalesce forms fail at the A2/A3 VPTO lowering boundary. - GM-to-L1 MGATHER is rejected explicitly at the A2/A3 VPTO lowering boundary until its separate NZ/scratch path is implemented. @@ -208,8 +214,8 @@ bisheng-unverified model (drops strides, no suffix, no bisheng test) — not use ### 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-authored via `pto.tile.gatherb` (block form), `pto.tile.gather` +(index form), and `pto.tile.mgather` (GM-to-UB Row/Elem), 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. diff --git a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp index f0e31a7176..622e58f28a 100644 --- a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp +++ b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp @@ -30,6 +30,7 @@ #include "llvm/ADT/SmallVector.h" #include +#include using namespace mlir; @@ -147,6 +148,8 @@ struct TileShapeMetadata { SmallVector validShape; pto::AddressSpace memorySpace; pto::BLayout bLayout; + pto::SLayout sLayout; + pto::CompactMode compactMode; }; using TileShapeMap = DenseMap; @@ -259,6 +262,7 @@ struct LowerPTOToUBufOpsPass for (auto op : allocOps) { auto tbTy = cast(op.getResult().getType()); auto shape = tbTy.getShape(); + auto config = tbTy.getConfigAttr(); Value addr = op.getAddr(); if (!addr) { op.emitError("A3 VPTO UB lowering requires planned alloc_tile " @@ -278,8 +282,10 @@ struct LowerPTOToUBufOpsPass SmallVector(tbTy.getValidShape()), cast(tbTy.getMemorySpace()) .getAddressSpace(), - tbTy.getConfigAttr() ? tbTy.getConfigAttr().getBLayout().getValue() - : pto::BLayout::RowMajor}; + config ? config.getBLayout().getValue() : pto::BLayout::RowMajor, + config ? config.getSLayout().getValue() : pto::SLayout::NoneBox, + config ? config.getCompactMode().getValue() + : pto::CompactMode::Null}; op.getResult().replaceAllUsesWith(pc.getResult()); op.erase(); } @@ -1434,8 +1440,54 @@ struct LowerPTOToUBufOpsPass SmallVector sizes; SmallVector strides; SmallVector offsets; + SmallVector staticSizes; + SmallVector staticStrides; }; + static int64_t getStaticValue(Value value) { + auto constant = getConstantIntValue(value); + return constant.value_or(ShapedType::kDynamic); + } + + static SmallVector extractStaticViewStrides(Value view) { + if (auto subview = view.getDefiningOp()) { + SmallVector sourceStrides = + extractStaticViewStrides(subview.getSource()); + ArrayRef subviewStrides = subview.getStaticStrides(); + if (sourceStrides.size() != subviewStrides.size()) + return SmallVector(subview.getType().getRank(), + ShapedType::kDynamic); + for (auto [sourceStride, subviewStride] : + llvm::zip_equal(sourceStrides, subviewStrides)) { + if (sourceStride == ShapedType::kDynamic || + subviewStride == ShapedType::kDynamic || sourceStride <= 0 || + subviewStride <= 0 || + sourceStride > + std::numeric_limits::max() / subviewStride) { + sourceStride = ShapedType::kDynamic; + continue; + } + sourceStride *= subviewStride; + } + return sourceStrides; + } + if (auto reinterpret = + view.getDefiningOp()) { + SmallVector strides; + for (OpFoldResult stride : reinterpret.getConstifiedMixedStrides()) { + auto value = getConstantIntValue(stride); + strides.push_back(value.value_or(ShapedType::kDynamic)); + } + return strides; + } + if (auto cast = view.getDefiningOp()) + return extractStaticViewStrides(cast.getSource()); + auto memTy = dyn_cast(view.getType()); + if (!memTy) + return {}; + return SmallVector(memTy.getStridesAndOffset().first); + } + static bool hasUnitInnermostStride(Value view) { while (Operation *def = view.getDefiningOp()) { if (auto subview = dyn_cast(def)) { @@ -1476,6 +1528,10 @@ struct LowerPTOToUBufOpsPass info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); info.strides.assign(mtvOp.getStrides().begin(), mtvOp.getStrides().end()); + for (Value size : info.sizes) + info.staticSizes.push_back(getStaticValue(size)); + for (Value stride : info.strides) + info.staticStrides.push_back(getStaticValue(stride)); if (info.strides.empty() || !matchPattern(info.strides.back(), m_One())) { op->emitError("A2/A3 DMA lowering requires a unit innermost stride"); @@ -1522,6 +1578,8 @@ struct LowerPTOToUBufOpsPass info.sizes.assign(metadata.getSizes().begin(), metadata.getSizes().end()); info.strides.assign(metadata.getStrides().begin(), metadata.getStrides().end()); + info.staticSizes.assign(shape.begin(), shape.end()); + info.staticStrides = extractStaticViewStrides(view); return info; } @@ -1738,20 +1796,6 @@ struct LowerPTOToUBufOpsPass llvm_unreachable("unknown MGATHER OOB mode"); } - Value flattenGatherIndex(Location loc, OpBuilder &b, Value flatIndex, - ArrayRef sizes, ArrayRef strides) { - Value remaining = flatIndex; - Value linearOffset = idxc0(loc, b); - for (int64_t i = static_cast(sizes.size()) - 1; i >= 0; --i) { - Value coordinate = b.create(loc, remaining, sizes[i]); - Value contribution = b.create(loc, coordinate, strides[i]); - linearOffset = b.create(loc, linearOffset, contribution); - if (i != 0) - remaining = b.create(loc, remaining, sizes[i]); - } - return linearOffset; - } - void emitScalarZero(Location loc, OpBuilder &b, Value dst, Value dstOffset, Type elemTy) { Value zero = @@ -1774,6 +1818,132 @@ struct LowerPTOToUBufOpsPass b.create(loc, pipeS, pipeMte2, event0); } + static bool hasPositiveStaticShape(ArrayRef shape) { + return !shape.empty() && llvm::all_of(shape, [](int64_t dim) { + return dim != ShapedType::kDynamic && dim > 0; + }); + } + + static bool isDenseRowMajor(ArrayRef shape, + ArrayRef strides) { + if (shape.size() != strides.size() || !hasPositiveStaticShape(shape)) + return false; + int64_t expectedStride = 1; + for (int64_t i = static_cast(shape.size()) - 1; i >= 0; --i) { + if (strides[i] != expectedStride) + return false; + if (shape[i] > std::numeric_limits::max() / expectedStride) + return false; + expectedStride *= shape[i]; + } + return true; + } + + static bool hasRepresentableByteExtent(ArrayRef shape, + ArrayRef strides, + unsigned elemSize) { + if (shape.size() != strides.size() || !hasPositiveStaticShape(shape) || + !hasPositiveStaticShape(strides) || elemSize == 0) + return false; + int64_t maxOffset = 0; + for (auto [dim, stride] : llvm::zip_equal(shape, strides)) { + int64_t count = dim - 1; + if (count > + (std::numeric_limits::max() - maxOffset) / stride) + return false; + maxOffset += count * stride; + } + return maxOffset <= + std::numeric_limits::max() / elemSize - 1; + } + + static bool isSupportedMGatherTile(const TileShapeMetadata &meta) { + return meta.memorySpace == pto::AddressSpace::VEC && + meta.bLayout == pto::BLayout::RowMajor && + meta.sLayout == pto::SLayout::NoneBox && + meta.compactMode != pto::CompactMode::RowPlusOne; + } + + LogicalResult validateMGatherSafeSubset( + pto::MGatherOp op, const DmaViewInfo &viewInfo, + const TileShapeMetadata &idxMeta, const TileShapeMetadata &dstMeta, + Type elemTy, pto::Coalesce coalesce) { + if (!isSupportedMGatherTile(dstMeta)) + return op.emitOpError( + "A2/A3 VPTO mgather requires a row-major, none-box, non-row_plus_one " + "VEC dst tile"); + if (!isSupportedMGatherTile(idxMeta)) + return op.emitOpError( + "A2/A3 VPTO mgather requires a row-major, none-box, non-row_plus_one " + "VEC idx tile"); + if (dstMeta.shape.size() != 2 || dstMeta.validShape.size() != 2 || + idxMeta.shape.size() != 2 || idxMeta.validShape.size() != 2) + return op.emitOpError("requires rank-2 idx and dst tiles"); + for (const TileShapeMetadata *meta : {&idxMeta, &dstMeta}) { + for (size_t i = 0; i < 2; ++i) { + if (meta->shape[i] <= 0 || meta->validShape[i] <= 0 || + meta->validShape[i] > meta->shape[i]) + return op.emitOpError( + "requires positive idx/dst valid shapes within physical shapes"); + } + } + + auto gmPtrTy = dyn_cast(viewInfo.gmPtr.getType()); + if (!gmPtrTy || gmPtrTy.getMemorySpace().getAddressSpace() != + pto::AddressSpace::GM) + return op.emitOpError("requires a GM source table"); + if (gmPtrTy.getElementType() != elemTy) + return op.emitOpError( + "requires the GM table element type to match dst exactly"); + if (!hasPositiveStaticShape(viewInfo.staticSizes) || + !hasPositiveStaticShape(viewInfo.staticStrides) || + viewInfo.staticSizes.size() != viewInfo.staticStrides.size()) + return op.emitOpError( + "requires statically known positive GM table dimensions and strides"); + + unsigned elemSize = getMGatherElementSize(elemTy); + if (!hasRepresentableByteExtent(viewInfo.staticSizes, + viewInfo.staticStrides, elemSize)) + return op.emitOpError( + "requires statically representable GM table byte addressing"); + if (coalesce == pto::Coalesce::Row) { + if (viewInfo.staticSizes.size() < 2) + return op.emitOpError( + "requires a rank-2 or greater GM table for coalesce=row"); + for (size_t i = 0; i + 2 < viewInfo.staticSizes.size(); ++i) { + if (viewInfo.staticSizes[i] != 1) + return op.emitOpError( + "supports coalesce=row only for an effective rank-2 ND table"); + } + if (viewInfo.staticStrides.back() != 1) + return op.emitOpError( + "requires a unit innermost GM stride for coalesce=row"); + if (viewInfo.staticSizes.back() != dstMeta.validShape[1]) + return op.emitOpError( + "requires the GM table row width to equal dst valid_col"); + if (viewInfo.staticStrides[viewInfo.staticStrides.size() - 2] < + viewInfo.staticSizes.back()) + return op.emitOpError( + "requires the GM table row stride to cover the logical row width"); + if (dstMeta.validShape[1] % (32 / elemSize) != 0) + return op.emitOpError( + "requires coalesce=row valid_col byte width to be 32-byte aligned"); + if (idxMeta.validShape[0] != 1 || + idxMeta.validShape[1] != dstMeta.validShape[0]) + return op.emitOpError( + "requires coalesce=row idx valid_shape to be [1, dst.valid_row]"); + return success(); + } + + if (idxMeta.validShape != dstMeta.validShape) + return op.emitOpError( + "requires coalesce=elem idx valid_shape to match dst valid_shape"); + if (!isDenseRowMajor(viewInfo.staticSizes, viewInfo.staticStrides)) + return op.emitOpError( + "requires a dense contiguous GM table for coalesce=elem"); + return success(); + } + LogicalResult emitMGatherRow(pto::MGatherOp op, OpBuilder &b, const DmaViewInfo &viewInfo, Value gmPtr, Value idxPtr, Value dstPtr, @@ -1785,20 +1955,16 @@ struct LowerPTOToUBufOpsPass return op.emitOpError( "requires a rank-2 or greater GM table for coalesce=row"); - ArrayRef rowSizes(viewInfo.sizes.data(), viewInfo.sizes.size() - 1); - ArrayRef rowStrides(viewInfo.strides.data(), - viewInfo.strides.size() - 1); - Value tableRows = product(loc, b, rowSizes); + size_t rowDim = viewInfo.sizes.size() - 2; + Value tableRows = viewInfo.sizes[rowDim]; + Value tableRowStride = viewInfo.strides[rowDim]; int64_t dstRows = dstMeta.validShape[0]; int64_t dstCols = dstMeta.validShape[1]; int64_t dstStride = dstMeta.shape[1]; - bool contiguousIndex = - idxMeta.validShape[0] <= 1 || idxMeta.bLayout == pto::BLayout::ColMajor; unsigned elemSize = getMGatherElementSize(elemTy); for (int64_t row = 0; row < dstRows; ++row) { - Value idxOffset = - idxc(contiguousIndex ? row : row * idxMeta.shape[1], loc, b); + Value idxOffset = idxc(row, loc, b); Value rawIndex = b.create(loc, b.getI32Type(), idxPtr, idxOffset); NormalizedGatherIndex index = @@ -1807,7 +1973,7 @@ struct LowerPTOToUBufOpsPass auto emitCopy = [&]() { Value sourceOffset = - flattenGatherIndex(loc, b, index.value, rowSizes, rowStrides); + b.create(loc, index.value, tableRowStride); Value source = b.create(loc, gmPtr.getType(), gmPtr, sourceOffset); Value destination = @@ -1865,11 +2031,8 @@ struct LowerPTOToUBufOpsPass b.create(loc, b.getI32Type(), idxPtr, idxOffset); NormalizedGatherIndex index = normalizeGatherIndex(loc, b, rawIndex, tableElements, oob); - Value sourceOffset = flattenGatherIndex( - loc, b, index.value, viewInfo.sizes, viewInfo.strides); - if (oob != pto::GatherOOB::Zero) { - emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, sourceOffset, + emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, index.value, elemTy); continue; } @@ -1878,7 +2041,7 @@ struct LowerPTOToUBufOpsPass /*addThenBlock=*/true, /*addElseBlock=*/true); b.setInsertionPointToStart(ifOp.thenBlock()); - emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, sourceOffset, + emitScalarGather(loc, b, gmPtr, dstPtr, dstOffset, index.value, elemTy); b.create(loc); b.setInsertionPointToStart(ifOp.elseBlock()); @@ -1907,17 +2070,16 @@ struct LowerPTOToUBufOpsPass if (op.getScratch()) return op.emitOpError( "A2/A3 GM->UB VPTO mgather does not support a scratch operand"); - if (dstIt->second.shape.size() != 2 || - dstIt->second.validShape.size() != 2 || - idxIt->second.shape.size() != 2 || idxIt->second.validShape.size() != 2) - return op.emitOpError("requires rank-2 idx and dst tiles"); Type elemTy = getStoredElemType(op.getDst().getType()); if (!elemTy || getMGatherElementSize(elemTy) == 0) return op.emitOpError("has an unsupported destination element type"); auto idxPtrTy = dyn_cast(op.getIdx().getType()); - if (!idxPtrTy || !idxPtrTy.getElementType().isInteger(32)) - return op.emitOpError("requires an i32 UB index tile"); + auto idxElemTy = idxPtrTy + ? dyn_cast(idxPtrTy.getElementType()) + : IntegerType{}; + if (!idxElemTy || idxElemTy.getWidth() != 32 || !idxElemTy.isSignless()) + return op.emitOpError("requires a signless i32 UB index tile"); auto viewInfo = extractDmaViewInfo(op.getOperation(), op.getMem()); if (failed(viewInfo)) @@ -1925,21 +2087,19 @@ struct LowerPTOToUBufOpsPass if (viewInfo->sizes.size() != viewInfo->strides.size()) return op.emitOpError("requires one stride per GM table dimension"); + auto coalesceAttr = + dyn_cast_or_null(op.getProperties().coalesce); + if (!coalesceAttr) + return op.emitOpError("requires an explicit coalesce attribute"); + pto::Coalesce coalesce = coalesceAttr.getValue(); + if (failed(validateMGatherSafeSubset(op, *viewInfo, idxIt->second, + dstIt->second, elemTy, coalesce))) + return failure(); + Value byteOff = computeGMByteOffset(op.getLoc(), b, *viewInfo, getMGatherElementSize(elemTy)); Value gmPtr = offsetGMPtrByBytes(op.getLoc(), b, viewInfo->gmPtr, byteOff); - auto coalesceAttr = - dyn_cast_or_null(op.getProperties().coalesce); - pto::Coalesce coalesce; - if (coalesceAttr) { - coalesce = coalesceAttr.getValue(); - } else if (idxIt->second.validShape == dstIt->second.validShape) { - coalesce = pto::Coalesce::Elem; - } else { - coalesce = pto::Coalesce::Row; - } - pto::GatherOOB oob = op.getGatherOob(); if (coalesce == pto::Coalesce::Row) return emitMGatherRow(op, b, *viewInfo, gmPtr, op.getIdx(), op.getDst(), diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index c3349f5459..52e8cd402e 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3789,6 +3789,32 @@ def tci(start, dst, *, tmp=None, descending=False): ) +def mgather(mem, idx, dst, coalesce, *, scratch=None, gather_oob=None): + """``pto.mgather`` from a GM tensor view into a destination tile.""" + if not isinstance(coalesce, Attribute): + coalesce = _enum_attr( + "coalesce", + coalesce, + supported={"row", "elem"}, + context="tile.mgather(..., coalesce=...)", + ) + if gather_oob is not None and not isinstance(gather_oob, Attribute): + gather_oob = _enum_attr( + "gather_oob", + gather_oob, + supported={"undefined", "clamp", "wrap", "zero"}, + context="tile.mgather(..., gather_oob=...)", + ) + _pto.mgather( + unwrap_surface_value(mem), + unwrap_surface_value(idx), + unwrap_surface_value(dst), + coalesce, + scratch=None if scratch is None else unwrap_surface_value(scratch), + gather_oob=gather_oob, + ) + + 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 0ceb3c534b..9543b2508c 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -179,6 +179,7 @@ def rowargmin(src, dst, *, tmp=None): mrgsort = staticmethod(_ops.tmrgsort) gather = staticmethod(_ops.tgather) gatherb = staticmethod(_ops.tgatherb) + mgather = staticmethod(_ops.mgather) tri = staticmethod(_ops.ttri) histogram = staticmethod(_ops.tthistogram) diff --git a/ptodsl/tests/e2e/common.py b/ptodsl/tests/e2e/common.py index d8a8096539..6e957a1f6c 100644 --- a/ptodsl/tests/e2e/common.py +++ b/ptodsl/tests/e2e/common.py @@ -716,3 +716,158 @@ def launch_and_check_gather_index( 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_mgather_kernel( + *, + src_rows: int, + src_cols: int, + dst_rows: int, + dst_cols: int, + coalesce: str, + gather_oob: str, + dtype_str: str = "float32", + target: str = "a3", + backend: str = "vpto", + kernel_kind: str = "vector", +): + """Return a ``@pto.jit`` KernelHandle for GM-to-UB ``pto.tile.mgather``.""" + pto_dtype = getattr(pto, dtype_str) + idx_dtype = pto.i32 + idx_rows = 1 if coalesce == "row" else dst_rows + idx_cols = dst_rows if coalesce == "row" else dst_cols + idx_physical_cols = ((idx_cols + 7) // 8) * 8 + fn_name = ( + f"mgather_{coalesce}_{gather_oob}_{dtype_str}_" + f"{src_rows}x{src_cols}_{dst_rows}x{dst_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_src_rows = pto.const(src_rows) + c_src_cols = pto.const(src_cols) + c_src_elems = pto.const(src_rows * src_cols) + c_idx_rows = pto.const(idx_rows) + c_idx_cols = pto.const(idx_physical_cols) + c_idx_elems = pto.const(idx_rows * idx_physical_cols) + c_dst_rows = pto.const(dst_rows) + c_dst_cols = pto.const(dst_cols) + c_dst_elems = pto.const(dst_rows * dst_cols) + + src_shape = [c1, c1, c1, c_src_rows, c_src_cols] + src_strides = [c_src_elems, c_src_elems, c_src_elems, c_src_cols, c1] + idx_shape = [c1, c1, c1, c_idx_rows, c_idx_cols] + idx_strides = [c_idx_elems, c_idx_elems, c_idx_elems, c_idx_cols, c1] + dst_shape = [c1, c1, c1, c_dst_rows, c_dst_cols] + dst_strides = [c_dst_elems, c_dst_elems, c_dst_elems, c_dst_cols, c1] + off = [c0, c0, c0, c0, c0] + + src_view = pto.make_tensor_view(Src_ptr, shape=src_shape, strides=src_strides) + idx_view = pto.make_tensor_view(Idx_ptr, shape=idx_shape, strides=idx_strides) + dst_view = pto.make_tensor_view(Dst_ptr, shape=dst_shape, strides=dst_strides) + src_part = pto.partition_view(src_view, offsets=off, sizes=src_shape) + idx_part = pto.partition_view(idx_view, offsets=off, sizes=idx_shape) + dst_part = pto.partition_view(dst_view, offsets=off, sizes=dst_shape) + + idx_tile = pto.alloc_tile( + shape=[idx_rows, idx_physical_cols], + valid_shape=[idx_rows, idx_cols], + dtype=idx_dtype, + ) + dst_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=pto_dtype) + pto.tile.load(idx_part, idx_tile) + pto.tile.mgather( + src_part, + idx_tile, + dst_tile, + coalesce, + gather_oob=gather_oob, + ) + 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_mgather( + *, + kernel_handle, + src_rows: int, + src_cols: int, + dst_rows: int, + dst_cols: int, + coalesce: str, + gather_oob: str, + dtype_str: str, + torch, + seed: int = 42, +): + """Compile, launch, and numerical-check one GM-to-UB mgather kernel.""" + rng = np.random.RandomState(seed) + np_dtype = np.float32 if dtype_str == "float32" else np.float16 + src = rng.randint(1, 10, size=(src_rows, src_cols)).astype(np_dtype) + idx_rows = 1 if coalesce == "row" else dst_rows + idx_cols = dst_rows if coalesce == "row" else dst_cols + idx_physical_cols = ((idx_cols + 7) // 8) * 8 + idx_count = idx_rows * idx_cols + limit = src_rows if coalesce == "row" else src_rows * src_cols + if gather_oob == "undefined": + idx = rng.randint(0, limit, size=(idx_count,), dtype=np.int32) + else: + pattern = np.array([-2, -1, 0, limit - 1, limit, limit + 1], dtype=np.int32) + idx = np.resize(pattern, idx_count) + + if gather_oob == "clamp": + normalized = np.clip(idx, 0, limit - 1) + elif gather_oob == "wrap": + normalized = idx % limit + else: + normalized = idx + + if coalesce == "row": + golden = np.zeros((dst_rows, dst_cols), dtype=np_dtype) + valid = (idx >= 0) & (idx < limit) + if gather_oob in {"clamp", "wrap"}: + valid[:] = True + golden[valid] = src[normalized[valid], :dst_cols] + else: + src_flat = src.reshape(-1) + golden_flat = np.zeros((idx_count,), dtype=np_dtype) + valid = (idx >= 0) & (idx < limit) + if gather_oob in {"clamp", "wrap"}: + valid[:] = True + golden_flat[valid] = src_flat[normalized[valid]] + golden = golden_flat.reshape(dst_rows, dst_cols) + + torch_dt = _torch_dtype(torch, dtype_str) + idx_storage = np.zeros((idx_rows, idx_physical_cols), dtype=np.int32) + if coalesce == "row": + idx_storage[0, :idx_cols] = idx + else: + idx_storage[:, :dst_cols] = idx.reshape(dst_rows, dst_cols) + src_dev = torch.from_numpy(src).to(device="npu:0", dtype=torch_dt) + idx_dev = torch.from_numpy(idx_storage).to(device="npu:0", dtype=torch.int32) + dst_dev = torch.empty((dst_rows, dst_cols), 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 + + np.testing.assert_array_equal(dst_dev.cpu().numpy(), golden) + return compile_s, launch_s diff --git a/ptodsl/tests/e2e/test_gather.py b/ptodsl/tests/e2e/test_gather.py index 2dc54525ce..625bcc4cf5 100644 --- a/ptodsl/tests/e2e/test_gather.py +++ b/ptodsl/tests/e2e/test_gather.py @@ -27,8 +27,10 @@ from .common import ( make_gather_index_kernel, make_gatherb_kernel, + make_mgather_kernel, launch_and_check_gather, launch_and_check_gather_index, + launch_and_check_mgather, ) @@ -87,6 +89,17 @@ for rows, cols, desc in GATHER_INDEX_SHAPES ] +MGATHER_PARAMS = [ + pytest.param(6, 16, 3, 16, "row", "undefined", "float16", id="mgather-row-undefined-f16"), + pytest.param(6, 8, 6, 8, "row", "zero", "float32", id="mgather-row-zero-f32"), + pytest.param(6, 16, 6, 16, "row", "clamp", "float16", id="mgather-row-clamp-f16"), + pytest.param(6, 8, 6, 8, "row", "wrap", "float32", id="mgather-row-wrap-f32"), + pytest.param(4, 8, 2, 8, "elem", "undefined", "float32", id="mgather-elem-undefined-f32"), + pytest.param(4, 16, 2, 16, "elem", "zero", "float16", id="mgather-elem-zero-f16"), + pytest.param(4, 8, 2, 8, "elem", "clamp", "float32", id="mgather-elem-clamp-f32"), + pytest.param(4, 16, 2, 16, "elem", "wrap", "float16", id="mgather-elem-wrap-f16"), +] + @pytest.mark.require_npu @pytest.mark.parametrize("case", GATHERB_F32_PARAMS) @@ -146,3 +159,49 @@ def test_gather_index(case, torch, target_arch, backend): ) print(f" PASS gather index {dtype_str} {rows}x{cols} ({desc}) " f"compile={compile_s:.3f}s launch={launch_s:.3f}s") + + +@pytest.mark.require_npu +@pytest.mark.parametrize( + "src_rows,src_cols,dst_rows,dst_cols,coalesce,gather_oob,dtype_str", + MGATHER_PARAMS, +) +def test_mgather( + src_rows, + src_cols, + dst_rows, + dst_cols, + coalesce, + gather_oob, + dtype_str, + torch, + target_arch, + backend, +): + kernel = make_mgather_kernel( + src_rows=src_rows, + src_cols=src_cols, + dst_rows=dst_rows, + dst_cols=dst_cols, + coalesce=coalesce, + gather_oob=gather_oob, + dtype_str=dtype_str, + target=target_arch, + backend=backend, + ) + compile_s, launch_s = launch_and_check_mgather( + kernel_handle=kernel, + src_rows=src_rows, + src_cols=src_cols, + dst_rows=dst_rows, + dst_cols=dst_cols, + coalesce=coalesce, + gather_oob=gather_oob, + dtype_str=dtype_str, + torch=torch, + ) + print( + f" PASS mgather {coalesce}/{gather_oob} {dtype_str} " + f"{src_rows}x{src_cols}->{dst_rows}x{dst_cols} " + f"compile={compile_s:.3f}s launch={launch_s:.3f}s" + ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 71b975ab34..768a44441d 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -639,6 +639,35 @@ def test_tile_sort_gather_wrappers_call_low_level_ops(self): with self.assertRaisesRegex(ValueError, "unsupported tile mask pattern"): pto.tile.gather(src, dst, mask_pattern="PAT_ALL") + coalesce_attr = object() + oob_attr = object() + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops, "_enum_attr", side_effect=[coalesce_attr, oob_attr]) as enum_attr, \ + patch.object(_ops._pto, "mgather") as mgather_op: + pto.tile.mgather(src, idx, dst, "row", gather_oob="zero") + self.assertEqual( + enum_attr.call_args_list, + [ + call( + "coalesce", + "row", + supported={"row", "elem"}, + context="tile.mgather(..., coalesce=...)", + ), + call( + "gather_oob", + "zero", + supported={"undefined", "clamp", "wrap", "zero"}, + context="tile.mgather(..., gather_oob=...)", + ), + ], + ) + self.assertEqual(mgather_op.call_args.args, (src, idx, dst, coalesce_attr)) + self.assertEqual( + mgather_op.call_args.kwargs, + {"scratch": None, "gather_oob": oob_attr}, + ) + def test_tile_mov_accepts_acc_to_vec_mode(self): src = object() dst = object() diff --git a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto index d32b934288..28c5315e72 100644 --- a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto +++ b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3.pto @@ -18,23 +18,25 @@ module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind %rowPart = pto.partition_view %rowView, offsets = [%c0, %c0], sizes = [%c4, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<4x16xf16> %rowIdx = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf %rowDst = pto.alloc_tile : !pto.tile_buf - pto.mgather ins(%rowPart, %rowIdx : !pto.partition_tensor_view<4x16xf16>, !pto.tile_buf) + pto.mgather ins(%rowPart, %rowIdx : !pto.partition_tensor_view<4x16xf16>, !pto.tile_buf) outs(%rowDst : !pto.tile_buf) {coalesce = #pto} // CHECK: pto.load_scalar + // CHECK: arith.muli {{.*}}, %strides{{.*}}#3 : index // CHECK: pto.set_flag[, , ] // CHECK: pto.wait_flag[, , ] - // CHECK: pto.mte_gm_ub + // CHECK: pto.mte_gm_ub {{.*}}, %c32{{.*}} nburst %zeroView = pto.make_tensor_view %zeroSrc, shape = [%c4, %c32], strides = [%c32, %c1] : !pto.tensor_view @@ -63,8 +65,9 @@ module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind, !pto.tile_buf) outs(%clampDst : !pto.tile_buf) {coalesce = #pto, gatherOob = #pto} - // CHECK: arith.select - // CHECK: pto.load_scalar {{.*}} : !pto.ptr -> f32 + // CHECK: %[[CLAMP_LOW:.*]] = arith.select + // CHECK: %[[CLAMP_IDX:.*]] = arith.select {{.*}}, {{.*}}, %[[CLAMP_LOW]] : index + // CHECK-NEXT: {{.*}} = pto.load_scalar {{.*}}[%[[CLAMP_IDX]]] : !pto.ptr -> f32 // CHECK: pto.store_scalar %wrapView = pto.make_tensor_view %wrapSrc, shape = [%c8], strides = [%c1] @@ -79,7 +82,8 @@ module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind) {coalesce = #pto, gatherOob = #pto} // CHECK: arith.remsi - // CHECK: pto.load_scalar {{.*}} : !pto.ptr -> i32 + // CHECK: %[[WRAP_IDX:.*]] = arith.select + // CHECK-NEXT: {{.*}} = pto.load_scalar {{.*}}[%[[WRAP_IDX]]] : !pto.ptr -> i32 // CHECK: pto.store_scalar // CHECK-NOT: pto.mgather return diff --git a/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_unsupported.pto b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_unsupported.pto new file mode 100644 index 0000000000..a203989fbc --- /dev/null +++ b/test/lit/vpto/ub/tile_to_ub/gather/mgather_a2a3_unsupported.pto @@ -0,0 +1,83 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto-llvm-ir %s -o /dev/null 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @row_col_major_idx(%mem: memref<4x16xf16, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: A2/A3 VPTO mgather requires a row-major, none-box, non-row_plus_one VEC idx tile + pto.mgather ins(%mem, %idx : memref<4x16xf16, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @row_unaligned_width(%mem: memref<4x15xf16, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: requires coalesce=row valid_col byte width to be 32-byte aligned + pto.mgather ins(%mem, %idx : memref<4x15xf16, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @row_width_mismatch(%mem: memref<4x16xf16, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: requires the GM table row width to equal dst valid_col + pto.mgather ins(%mem, %idx : memref<4x16xf16, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @elem_empty_table(%mem: memref<0x8xf32, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: requires statically known positive GM table dimensions and strides + pto.mgather ins(%mem, %idx : memref<0x8xf32, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @elem_non_dense(%mem: memref<4x8xf32, strided<[16, 1]>, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: requires a dense contiguous GM table for coalesce=elem + pto.mgather ins(%mem, %idx : memref<4x8xf32, strided<[16, 1]>, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @elem_unrepresentable_extent(%mem: memref<2x8xf32, strided<[9223372036854775807, 1]>, #pto.address_space>) attributes {pto.aicore} { + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + // CHECK: requires statically representable GM table byte addressing + pto.mgather ins(%mem, %idx : memref<2x8xf32, strided<[9223372036854775807, 1]>, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } +}