From 04daa2d60a6ba37cfa7611484027b90f0260feaf Mon Sep 17 00:00:00 2001 From: likai00 Date: Sat, 29 Aug 2026 14:58:32 +0800 Subject: [PATCH] vgather support i8/ui8 -> i16/ui16 --- docs/isa/vmi-isa/07-sfu.md | 24 +- include/PTO/IR/VMIOps.td | 13 +- include/PTO/IR/VPTOOps.td | 6 + lib/PTO/IR/VMI.cpp | 95 +++- lib/PTO/Transforms/VMIToVPTO.cpp | 70 ++- lib/PTO/Transforms/VPTOMaskSimplify.cpp | 18 +- ptodsl/ptodsl/_vmi_namespace.py | 15 +- .../vmi_to_vpto_gather_all_active_mask.pto | 37 ++ ...pto_gather_all_active_mask_multi_chunk.pto | 43 ++ ...vmi_to_vpto_gather_chunk_sweep_i32_i32.pto | 115 ++++ .../vmi_to_vpto_gather_chunk_sweep_i8_i16.pto | 158 ++++++ ...vmi_to_vpto_gather_chunk_sweep_u16_u16.pto | 158 ++++++ ...mi_to_vpto_gather_chunk_sweep_ui8_ui16.pto | 158 ++++++ .../vmi_new/vmi_to_vpto_gather_f16_bf16.pto | 57 ++ .../vmi_to_vpto_gather_f16_invalid.pto | 4 +- ...mi_to_vpto_gather_granularity_conflict.pto | 38 ++ ...mi_to_vpto_gather_granularity_explicit.pto | 31 ++ .../vmi_to_vpto_gather_granularity_pred.pto | 26 + test/lit/vmi_new/vmi_to_vpto_gather_i16.pto | 36 ++ .../vmi_new/vmi_to_vpto_gather_i8_to_i16.pto | 57 ++ ...i_to_vpto_gather_i8_to_i16_multi_chunk.pto | 62 +++ ...i_to_vpto_gather_scatter_shape_invalid.pto | 4 +- ...pto_gather_too_many_chunks_invalid_i32.pto | 27 + ...to_gather_too_many_chunks_invalid_ui16.pto | 27 + test/lit/vpto/vpto_mask_simplify.pto | 57 ++ .../cases/vmi_new/vgather-dsl-i8-ui8-256.py | 166 ++++++ test/vpto/cases/vmi_new/vgather-vmi.py | 516 ++++++++++++++++++ 27 files changed, 1973 insertions(+), 45 deletions(-) create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask_multi_chunk.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i32_i32.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i8_i16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_u16_u16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_ui8_ui16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_f16_bf16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_granularity_conflict.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_granularity_explicit.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_granularity_pred.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_i16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16_multi_chunk.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_i32.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_ui16.pto create mode 100644 test/vpto/cases/vmi_new/vgather-dsl-i8-ui8-256.py create mode 100644 test/vpto/cases/vmi_new/vgather-vmi.py diff --git a/docs/isa/vmi-isa/07-sfu.md b/docs/isa/vmi-isa/07-sfu.md index 6b73faa81f..e69f63a324 100644 --- a/docs/isa/vmi-isa/07-sfu.md +++ b/docs/isa/vmi-isa/07-sfu.md @@ -413,7 +413,7 @@ ### `pto.vmi.vgather` -- **semantics:** Indexed gather from UB at B32 granularity. For each active +- **semantics:** Indexed gather from UB at B32/B16 granularity. For each active lane `i`, load `src[offsets[i]]`. ```c @@ -423,7 +423,17 @@ - **syntax:** ```mlir - %g = pto.vmi.vgather %src, %offsets, %mask : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + // B32 path + %g = pto.vmi.vgather %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg // T in {i32,ui32,f32} + + // B16 path + %g = pto.vmi.vgather %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg // T in {i16,ui16,f16,bf16} + %g = pto.vmi.vgather %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg + %g = pto.vmi.vgather %src, %offsets, %mask + : !pto.ptr, !pto.vmi.vreg, !pto.vmi.mask -> !pto.vmi.vreg ``` - **operands:** @@ -435,12 +445,10 @@ - **results:** `!pto.vmi.vreg` - **attributes:** `pmode` -- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32` -- **lowering to `pto.mi`:** - ``` - K × pto.vgather2 - ``` - `#mi = K`, `dep = 1`, util data-dependent. +- **datatypes:** B32 -- `i32`/`ui32`/`f32`; B16 -- `i16`/`ui16`/`f16`/`bf16`, + plus `i8`/`ui8` -> `i16`/`ui16` zero-extension. +- **lowering:** B16 -> `K × pto.vgather2`; B32 -> `K × pto.vgather2_bc`. + A statically all-active mask omits the trailing `vsel`. ### `pto.vmi.vgatherb` diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 110a5bf95d..a2aaecb5a7 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -710,6 +710,13 @@ def VMIMaskedLoadOp : VMI_Op<"masked_load", [DeclareOpInterfaceMethods]> { let summary = "VMI logical masked indexed gather with passthrough lanes"; + let description = [{ + Gathers elements from UB memory using an index vector. + i8/ui8 sources may be gathered into matching i16/ui16 results. The + 8-to-16 promotion is a zero-extension: one byte is read per lane and the + upper 8 bits are zero-filled, regardless of whether the source element + type is signed (i8) or unsigned (ui8). Sign-extension is not supported. + }]; let arguments = (ins PtrOrMemRef:$source, VMI_VRegTypeConstraint:$indices, VMI_MaskTypeConstraint:$mask, @@ -1584,9 +1591,13 @@ def VMIVdintlvOp : VMI_Op<"vdintlv"> { } def VMIVgatherOp : VMI_Op<"vgather", [DeclareOpInterfaceMethods]> { - let summary = "VMI logical masked indexed gather (B32 granularity)"; + let summary = "VMI logical masked indexed gather"; let description = [{ Gathers elements from UB memory using an index vector. + i8/ui8 sources may be gathered into matching i16/ui16 results. The + 8-to-16 promotion is a zero-extension: one byte is read per lane and the + upper 8 bits are zero-filled, regardless of whether the source element + type is signed (i8) or unsigned (ui8). Sign-extension is not supported. Inactive lanes are controlled by pmode (no explicit passthru operand). pmode controls inactive-lane behavior: diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index a5cdd30843..6b5b6c5e18 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -2717,6 +2717,12 @@ def PTO_Vmrgsort4Op : PTO_VectorMicroOp<"vmrgsort4"> { def PTO_Vgather2Op : PTO_VectorMicroOp<"vgather2", [ DeclareOpInterfaceMethods ]> { + let description = [{ + Masked indexed gather from a UB-backed source. For 8-bit integer sources + the result is a 16-bit integer (i8/ui8 -> i16/ui16) produced by + zero-extension (one byte per lane, upper 8 bits zero-filled). + Sign-extension is not performed. + }]; let arguments = (ins PTO_BufferType:$source, PTO_VectorType:$offsets, diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index 0f42970e79..31bd858239 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -562,6 +562,72 @@ static LogicalResult verifyMemoryElementMatches(Operation *op, Type memoryType, return success(); } +// 8->16 gather promotion is a zero-extension (unsigned) operation. signless +// i8/i16 are accepted and treated as unsigned bytes; sign-extension is not +// supported (see VMIVgatherOp / Vgather2Op description). +static bool isVMI8To16GatherPair(Type sourceElemType, Type resultElemType) { + auto srcInt = dyn_cast(sourceElemType); + auto resInt = dyn_cast(resultElemType); + if (!srcInt || !resInt || + srcInt.getWidth() != mlir::pto::kValue8 || + resInt.getWidth() != mlir::pto::kValue16) { + return false; + } + if (srcInt.isUnsigned()) { + return resInt.isUnsigned(); + } + return !resInt.isUnsigned(); +} + +static LogicalResult verifyGatherMemoryElementMatches( + Operation *op, Type memoryType, VMIVRegType dataType, StringRef role) { + Type memoryElementType = getMemoryElementType(memoryType); + if (!memoryElementType) { + return success(); + } + if (memoryElementType == dataType.getElementType()) { + return success(); + } + if (isVMI8To16GatherPair(memoryElementType, dataType.getElementType())) { + return success(); + } + return op->emitOpError() + << "requires memory " << role + << " element type to match VMI data element type" + " or be an 8-bit integer promoted to a matching 16-bit integer"; +} + +static bool isSameWidth16BitGatherPair(Type sourceElemType, + Type resultElemType) { + // Existing VMI f16/bf16 path: same-width 16-bit float gather. + if (sourceElemType == resultElemType && + (sourceElemType.isF16() || sourceElemType.isBF16())) { + return true; + } + auto srcInt = dyn_cast(sourceElemType); + auto resInt = dyn_cast(resultElemType); + // Existing VMI ui16/i16 path: same-width 16-bit integer gather with matching + // integer semantics (signless i16 / i16 is accepted as the non-unsigned side). + if (!srcInt || !resInt || + srcInt.getWidth() != mlir::pto::kValue16 || + resInt.getWidth() != mlir::pto::kValue16) { + return false; + } + if (srcInt.isUnsigned()) { + return resInt.isUnsigned(); + } + return !resInt.isUnsigned(); +} + +static bool isSupported16BitGatherResult(Type sourceElemType, + Type resultElemType) { + // New 8 -> 16 path: i8/ui8 -> i16/ui16 with matching integer semantics. + if (isVMI8To16GatherPair(sourceElemType, resultElemType)) { + return true; + } + return isSameWidth16BitGatherPair(sourceElemType, resultElemType); +} + static LogicalResult verifyContiguousIfLayoutAssigned(Operation *op, VMIVRegType type, StringRef role) { @@ -2582,7 +2648,7 @@ LogicalResult VMIGatherOp::verify() { auto maskType = cast(getMask().getType()); auto passthruType = cast(getPassthru().getType()); auto resultType = cast(getResult().getType()); - if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + if (failed(verifyGatherMemoryElementMatches(getOperation(), getSource().getType(), resultType, "source"))) { return failure(); } @@ -2609,13 +2675,13 @@ LogicalResult VMIGatherOp::verify() { return failure(); } - auto resultIntegerType = dyn_cast(resultType.getElementType()); if (indexElementType.getWidth() == mlir::pto::kValue16 && - (!resultIntegerType || !resultIntegerType.isUnsigned() || - resultIntegerType.getWidth() != mlir::pto::kValue16)) { + !isSupported16BitGatherResult( + getMemoryElementType(getSource().getType()), + resultType.getElementType())) { return emitOpError( - "requires ui16 result and passthru element type when using ui16 " - "indices"); + "requires i16/ui16/f16/bf16 result and passthru element type when " + "using ui16 indices, or i8/ui8 -> i16/ui16 integer promotion"); } return verifyMaskMatchesData(getOperation(), maskType, resultType); } @@ -4022,7 +4088,7 @@ LogicalResult VMIVgatherOp::verify() { return failure(); } - if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(), + if (failed(verifyGatherMemoryElementMatches(getOperation(), getSource().getType(), resultType, "source"))) { return failure(); } @@ -4044,16 +4110,15 @@ LogicalResult VMIVgatherOp::verify() { return failure(); } - // 16-bit offsets only address the ui16 gather path (pto.vgather2 / b16 mask), - // which requires a ui16 result element type. Reject other 16-bit-offset - // results here so the error surfaces at the vgather op rather than later in - // the legacy gather it lowers to. - auto resultIntegerType = dyn_cast(resultType.getElementType()); + // 16-bit offsets address the pto.vgather2 / b16 mask path. It supports + // i16/ui16/f16/bf16 same-width gather and i8/ui8 -> i16/ui16 promotion. if (indexElementType.getWidth() == mlir::pto::kValue16 && - (!resultIntegerType || !resultIntegerType.isUnsigned() || - resultIntegerType.getWidth() != mlir::pto::kValue16)) { + !isSupported16BitGatherResult( + getMemoryElementType(getSource().getType()), + resultType.getElementType())) { return emitOpError( - "requires ui16 result element type when using ui16 offsets"); + "requires i16/ui16/f16/bf16 result element type when using ui16 " + "offsets, or i8/ui8 -> i16/ui16 integer promotion"); } if (auto pmode = getPmode()) { diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 8a1208c6c1..280a50842c 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -1941,14 +1941,37 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { auto indexElementType = dyn_cast(indicesType.getElementType()); if (!indexElementType || indexElementType.isSigned()) return fail("requires signless or unsigned integer indices"); - bool isU16Gather = resultBits == 16 && indexElementType.isUnsigned() && - indexElementType.getWidth() == 16 && - maskType.getGranularity() == "b16"; + Type sourceElemType = getMemoryElementType(op.getSource().getType()); + auto sourceInt = dyn_cast(sourceElemType); + auto resultInt = dyn_cast(resultType.getElementType()); + bool isB8To16Gather = + resultBits == 16 && sourceInt && resultInt && + sourceInt.getWidth() == mlir::pto::kValue8 && + resultInt.getWidth() == mlir::pto::kValue16 && + indexElementType.isUnsigned() && + indexElementType.getWidth() == 16 && + maskType.getGranularity() == "b16" && + ((sourceInt.isUnsigned() && resultInt.isUnsigned()) || + (!sourceInt.isUnsigned() && !resultInt.isUnsigned())); + bool isSameWidth16Gather = + resultBits == 16 && indexElementType.isUnsigned() && + indexElementType.getWidth() == 16 && + maskType.getGranularity() == "b16" && + ((sourceInt && resultInt && + sourceInt.getWidth() == mlir::pto::kValue16 && + resultInt.getWidth() == mlir::pto::kValue16 && + ((sourceInt.isUnsigned() && resultInt.isUnsigned()) || + (!sourceInt.isUnsigned() && !resultInt.isUnsigned()))) || + (sourceElemType == resultType.getElementType() && + (sourceElemType.isF16() || sourceElemType.isBF16()))); + bool isB16Gather = isSameWidth16Gather || isB8To16Gather; bool isB32Gather = resultBits == 32 && indexElementType.getWidth() == 32 && maskType.getGranularity() == "b32"; - if (!isU16Gather && !isB32Gather) + if (!isB16Gather && !isB32Gather) { return fail("requires either 32-bit results with 32-bit indices and b32 " - "mask, or ui16 results with ui16 indices and b16 mask"); + "mask, or ui16/i16/f16/bf16 results with ui16 indices and " + "b16 mask (including i8/ui8 -> i16/ui16 promotion)"); + } FailureOr resultArity = getVMIPhysicalArity(resultType); FailureOr indicesArity = getVMIPhysicalArity(indicesType); @@ -1962,7 +1985,16 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { return fail("requires result, indices, passthru, and mask to have the " "same physical arity"); - if (isB32Gather) { + // Each pto.vgather2/pto.vgather2_bc emits one physical vector register per + // result part. The ISA has a hard limit of four physical registers per + // pto.vmi instruction, so reject anything above that instead of silently + // lowering to five or more physical gathers. + if (*resultArity > mlir::pto::kValue4) { + return fail("gather exceeds the 4 physical register limit per VMI " + "instruction"); + } + + if (isB32Gather || (isB16Gather && *resultArity != 1)) { std::string resultReason; std::string indicesReason; std::string passthruReason; @@ -1978,8 +2010,6 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) { passthruReason); if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason))) return fail(Twine("mask requires full physical chunks; ") + maskReason); - } else if (*resultArity != 1) { - return fail("ui16 gather currently supports one physical chunk"); } return success(); @@ -7048,6 +7078,13 @@ struct OneToNVMIGatherOpPattern : OneToNOpConversionPattern { indicesParts.size() != passthruParts.size() || indicesParts.size() != resultTypes.size()) return rewriter.notifyMatchFailure(op, "gather physical arity mismatch"); + // Static all-active masks select gathered[0] for every lane, so the + // trailing vsel is a semantic no-op. Skip it and keep gathered directly. + // Non-static masks still take the original gather + vsel path. + auto resultVMIType = cast(op.getResult().getType()); + bool allActive = isStaticAllActiveMask(op.getMask(), + resultVMIType.getElementCount()); + SmallVector results; results.reserve(resultTypes.size()); @@ -7069,10 +7106,14 @@ struct OneToNVMIGatherOpPattern : OneToNOpConversionPattern { .create(op.getLoc(), resultType, *source, indices, mask) .getResult(); - results.push_back( - rewriter - .create(op.getLoc(), resultType, gathered, passthru, mask) - .getResult()); + if (allActive) { + results.push_back(gathered); + } else { + results.push_back( + rewriter + .create(op.getLoc(), resultType, gathered, passthru, mask) + .getResult()); + } } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); @@ -13394,9 +13435,10 @@ verifySupportedVMIToVPTOOps(ModuleOp module, return WalkResult::advance(); gather.emitError() << kVMIDiagUnsupportedPrefix - << "pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only " + << "pto.vmi.gather lowers through pto.vgather2/pto.vgather2_bc + pto.vsel only " "for UB pointer sources, contiguous full physical chunks, " - "32-bit result elements, i32 indices, and b32 masks (" + "ui16/i16/f16/bf16 results with ui16 indices and b16 masks, " + "or 32-bit results with i32 indices and b32 masks (" << reason << ")"; return WalkResult::interrupt(); } diff --git a/lib/PTO/Transforms/VPTOMaskSimplify.cpp b/lib/PTO/Transforms/VPTOMaskSimplify.cpp index bbc0f8a727..e0691e0975 100644 --- a/lib/PTO/Transforms/VPTOMaskSimplify.cpp +++ b/lib/PTO/Transforms/VPTOMaskSimplify.cpp @@ -52,11 +52,27 @@ struct SimplifyAllTruePredicateReorder : public OpRewritePattern { } }; +struct SimplifyVselAllTrueMask : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(VselOp op, + PatternRewriter &rewriter) const override { + // vsel(src0, src1, mask) selects src0 for all-true masks. + if (!isAllTrueMask(op.getMask())) { + return failure(); + } + + rewriter.replaceOp(op, op.getSrc0()); + return success(); + } +}; + struct VPTOMaskSimplifyPass : public pto::impl::VPTOMaskSimplifyBase { void runOnOperation() override { RewritePatternSet patterns(&getContext()); - patterns.add, + patterns.add, SimplifyAllTruePredicateReorder, SimplifyAllTruePredicateReorder, SimplifyAllTruePredicateReorder, diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 457bea1408..f9b49d759c 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -432,10 +432,21 @@ def _derive_hist_result_type(acc, *, context: str): def _derive_vgather_result_type(source, offsets, *, context: str): offsets_type = _as_vmi_vreg_type(_type_of(offsets), context=context) - result_type = _pointer_element_type(_type_of(source), context=context) + result_elem_type = _pointer_element_type(_type_of(source), context=context) + # 8-bit integer sources use the B16 gather promotion path (i8 -> i16, + # ui8 -> ui16, si8 -> si16). + if IntegerType.isinstance(result_elem_type): + source_int_type = IntegerType(result_elem_type) + if source_int_type.width == 8: + if source_int_type.is_unsigned: + result_elem_type = IntegerType.get_unsigned(16) + elif source_int_type.is_signed: + result_elem_type = IntegerType.get_signed(16) + else: + result_elem_type = IntegerType.get_signless(16) return _pto.VMIVRegType.get( offsets_type.element_count, - result_type, + result_elem_type, layout=offsets_type.layout, ) diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask.pto b/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask.pto new file mode 100644 index 0000000000..8d68a5d248 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask.pto @@ -0,0 +1,37 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s + +// A statically all-active mask selects gathered[0] for every lane, so +// collect_gather must lower to a bare vgather2 without the trailing vsel. + +module { + func.func @gather_all_active_mask_u16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> { + %c32 = arith.constant 32 : index + %mask = pto.vmi.create_mask %c32 : index -> !pto.vmi.mask<32xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + return %out : !pto.vmi.vreg<32xui16, #pto.vmi.layout> + } +} + +// CHECK-LABEL: func.func @gather_all_active_mask_u16( +// CHECK: %[[GATHER:.*]] = pto.vgather2 +// CHECK-SAME: !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xui16> +// CHECK-NOT: pto.vsel +// CHECK: return %[[GATHER]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask_multi_chunk.pto b/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask_multi_chunk.pto new file mode 100644 index 0000000000..91386e8c3d --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_all_active_mask_multi_chunk.pto @@ -0,0 +1,43 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// A statically all-active 512-lane mask selects gathered[0] for every part. +// The all-active shortcut must apply per part in the multi-chunk loop and +// therefore must emit four bare vgather2 ops without trailing vsel. + +module { + func.func @gather_all_active_mask_i8_to_i16_512( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) { + %c512 = arith.constant 512 : index + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xi16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) + return %p0, %p1, %p2, %p3 + : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16> + } +} + +// CHECK-LABEL: func.func @gather_all_active_mask_i8_to_i16_512( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i32_i32.pto b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i32_i32.pto new file mode 100644 index 0000000000..ed2dedbd80 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i32_i32.pto @@ -0,0 +1,115 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Chunk sweep for physical block counts {1,2,4} x {static all-active, dynamic mask}. +// C=4 is the ISA limit of four physical registers per VMI instruction. + +module { + func.func @gather_i32_i32_2chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %c128 = arith.constant 128 : index + %mask = pto.vmi.create_mask %c128 : index -> !pto.vmi.mask<128xb32, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<128xi32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + } + + func.func @gather_i32_i32_2chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<128xi32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<128xb32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<128xi32, #pto.vmi.layout>, + !pto.vmi.mask<128xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + } + + func.func @gather_i32_i32_4chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %c256 = arith.constant 256 : index + %mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xb32, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xi32, #pto.vmi.layout>, + !pto.vmi.mask<256xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<256xi32, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1, %p2, %p3 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32> + } + + func.func @gather_i32_i32_4chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xi32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<256xb32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xi32, #pto.vmi.layout>, + !pto.vmi.mask<256xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<256xi32, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1, %p2, %p3 : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.vreg<64xi32> + } +} + +// CHECK-LABEL: func.func @gather_i32_i32_2chunk_allactive( +// CHECK: pto.vgather2_bc +// CHECK: pto.vgather2_bc +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i32_i32_2chunk_dynamic( +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i32_i32_4chunk_allactive( +// CHECK: pto.vgather2_bc +// CHECK: pto.vgather2_bc +// CHECK: pto.vgather2_bc +// CHECK: pto.vgather2_bc +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i32_i32_4chunk_dynamic( +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: pto.vgather2_bc +// CHECK: pto.vsel +// CHECK: return +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i8_i16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i8_i16.pto new file mode 100644 index 0000000000..ff8f9cbe20 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_i8_i16.pto @@ -0,0 +1,158 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Chunk sweep for physical block counts {1,2,4} x {static all-active, dynamic mask}. +// C=4 is the ISA limit of four physical registers per VMI instruction. + +module { + func.func @gather_i8_i16_1chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>) { + %c32 = arith.constant 32 : index + %mask = pto.vmi.create_mask %c32 : index -> !pto.vmi.mask<32xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xi16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>) + return %p0 : !pto.vreg<128xi16> + } + + func.func @gather_i8_i16_1chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xi16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>) + return %p0 : !pto.vreg<128xi16> + } + + func.func @gather_i8_i16_2chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>) { + %c256 = arith.constant 256 : index + %mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xi16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>) + return %p0, %p1 : !pto.vreg<128xi16>, !pto.vreg<128xi16> + } + + func.func @gather_i8_i16_2chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<256xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xi16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>) + return %p0, %p1 : !pto.vreg<128xi16>, !pto.vreg<128xi16> + } + + func.func @gather_i8_i16_4chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) { + %c512 = arith.constant 512 : index + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xi16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16> + } + + func.func @gather_i8_i16_4chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<512xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xi16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xi16, #pto.vmi.layout>) + -> (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.vreg<128xi16> + } +} + +// CHECK-LABEL: func.func @gather_i8_i16_1chunk_allactive( +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i8_i16_1chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i8_i16_2chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i8_i16_2chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i8_i16_4chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_i8_i16_4chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_u16_u16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_u16_u16.pto new file mode 100644 index 0000000000..25a76cc02d --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_u16_u16.pto @@ -0,0 +1,158 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Chunk sweep for physical block counts {1,2,4} x {static all-active, dynamic mask}. +// C=4 is the ISA limit of four physical registers per VMI instruction. + +module { + func.func @gather_u16_u16_1chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) { + %c32 = arith.constant 32 : index + %mask = pto.vmi.create_mask %c32 : index -> !pto.vmi.mask<32xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) + return %p0 : !pto.vreg<128xui16> + } + + func.func @gather_u16_u16_1chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) + return %p0 : !pto.vreg<128xui16> + } + + func.func @gather_u16_u16_2chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %c256 = arith.constant 256 : index + %mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xui16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1 : !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_u16_u16_2chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<256xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xui16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1 : !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_u16_u16_4chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %c512 = arith.constant 512 : index + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xui16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_u16_u16_4chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<512xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xui16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16> + } +} + +// CHECK-LABEL: func.func @gather_u16_u16_1chunk_allactive( +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_u16_u16_1chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_u16_u16_2chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_u16_u16_2chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_u16_u16_4chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_u16_u16_4chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_ui8_ui16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_ui8_ui16.pto new file mode 100644 index 0000000000..c0b9aa27e7 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_chunk_sweep_ui8_ui16.pto @@ -0,0 +1,158 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Chunk sweep for physical block counts {1,2,4} x {static all-active, dynamic mask}. +// C=4 is the ISA limit of four physical registers per VMI instruction. + +module { + func.func @gather_ui8_ui16_1chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) { + %c32 = arith.constant 32 : index + %mask = pto.vmi.create_mask %c32 : index -> !pto.vmi.mask<32xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) + return %p0 : !pto.vreg<128xui16> + } + + func.func @gather_ui8_ui16_1chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + %p0 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>) + return %p0 : !pto.vreg<128xui16> + } + + func.func @gather_ui8_ui16_2chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %c256 = arith.constant 256 : index + %mask = pto.vmi.create_mask %c256 : index -> !pto.vmi.mask<256xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xui16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1 : !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_ui8_ui16_2chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<256xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xui16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1 : !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_ui8_ui16_4chunk_allactive( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %c512 = arith.constant 512 : index + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xb16, #pto.vmi.layout> + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xui16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @gather_ui8_ui16_4chunk_dynamic( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<512xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xui16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1, %p2, %p3 : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16> + } +} + +// CHECK-LABEL: func.func @gather_ui8_ui16_1chunk_allactive( +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_ui8_ui16_1chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_ui8_ui16_2chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_ui8_ui16_2chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_ui8_ui16_4chunk_allactive( +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vsel +// CHECK: return + +// CHECK-LABEL: func.func @gather_ui8_ui16_4chunk_dynamic( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: return +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_f16_bf16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_f16_bf16.pto new file mode 100644 index 0000000000..02d8af9e61 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_f16_bf16.pto @@ -0,0 +1,57 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +module { + func.func @vmi_to_vpto_gather_f16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vreg<128xf16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xf16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xf16, #pto.vmi.layout>) + -> !pto.vreg<128xf16> + return %part : !pto.vreg<128xf16> + } + + func.func @vmi_to_vpto_gather_bf16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vreg<128xbf16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xbf16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xbf16, #pto.vmi.layout>) + -> !pto.vreg<128xbf16> + return %part : !pto.vreg<128xbf16> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_f16( +// CHECK: %[[GATHER_F16:.*]] = pto.vgather2 %arg0, %arg1, %arg2 : !pto.ptr, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xf16> +// CHECK: %[[OUT_F16:.*]] = pto.vsel %[[GATHER_F16]], %{{.*}}, %arg2 : !pto.vreg<128xf16>, !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> +// CHECK: return %[[OUT_F16]] + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_bf16( +// CHECK: %[[GATHER_BF16:.*]] = pto.vgather2 %arg0, %arg1, %arg2 : !pto.ptr, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xbf16> +// CHECK: %[[OUT_BF16:.*]] = pto.vsel %[[GATHER_BF16]], %{{.*}}, %arg2 : !pto.vreg<128xbf16>, !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<128xbf16> +// CHECK: return %[[OUT_BF16]] + +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_f16_invalid.pto b/test/lit/vmi_new/vmi_to_vpto_gather_f16_invalid.pto index 5042859e71..18ae66125e 100644 --- a/test/lit/vmi_new/vmi_to_vpto_gather_f16_invalid.pto +++ b/test/lit/vmi_new/vmi_to_vpto_gather_f16_invalid.pto @@ -22,5 +22,5 @@ module { } } -// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only -// CHECK-SAME: 32-bit result elements +// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2/pto.vgather2_bc + pto.vsel only +// CHECK-SAME: or 32-bit results diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_granularity_conflict.pto b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_conflict.pto new file mode 100644 index 0000000000..0987fc803a --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_conflict.pto @@ -0,0 +1,38 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment | FileCheck %s --check-prefix=GRAN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s --check-prefix=VPTO + +module { + func.func @gather_pred_conflict_ui16_i32( + %src16: !pto.ptr, + %idx16: !pto.vmi.vreg<128xui16>, + %src32: !pto.ptr, + %idx32: !pto.vmi.vreg<128xi32>, + %mask: !pto.vmi.mask<128xpred>) + -> (!pto.vmi.vreg<128xui16>, !pto.vmi.vreg<128xi32>) { + %g16 = pto.vmi.vgather %src16, %idx16, %mask + : !pto.ptr, !pto.vmi.vreg<128xui16>, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xui16> + %g32 = pto.vmi.vgather %src32, %idx32, %mask + : !pto.ptr, !pto.vmi.vreg<128xi32>, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xi32> + return %g16, %g32 : !pto.vmi.vreg<128xui16>, !pto.vmi.vreg<128xi32> + } +} + +// GRAN-LABEL: func.func @gather_pred_conflict_ui16_i32( +// GRAN: pto.vmi.ensure_mask_granularity + +// VPTO-LABEL: func.func @gather_pred_conflict_ui16_i32( +// VPTO: pto.vgather2 +// VPTO: pto.vgather2_bc +// VPTO-NOT: pto.vmi. +// VPTO-NOT: !pto.vmi. +// VPTO-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_granularity_explicit.pto b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_explicit.pto new file mode 100644 index 0000000000..4771d16de5 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_explicit.pto @@ -0,0 +1,31 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s + +module { + func.func @gather_explicit_b16_respected( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + return %out : !pto.vmi.vreg<32xui16, #pto.vmi.layout> + } +} + +// CHECK-LABEL: func.func @gather_explicit_b16_respected( +// CHECK: pto.vgather2 +// CHECK-NOT: pto.vmi.ensure_mask_granularity +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_granularity_pred.pto b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_pred.pto new file mode 100644 index 0000000000..b8d6570f39 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_granularity_pred.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 root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment | FileCheck %s + +module { + func.func @gather_pred_inferred_b16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16>, + %mask: !pto.vmi.mask<32xpred>) + -> !pto.vmi.vreg<32xui16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, !pto.vmi.vreg<32xui16>, !pto.vmi.mask<32xpred> + -> !pto.vmi.vreg<32xui16> + return %out : !pto.vmi.vreg<32xui16> + } +} + +// CHECK-LABEL: func.func @gather_pred_inferred_b16( +// CHECK: !pto.vmi.mask<32xb16> +// CHECK-NOT: !pto.vmi.mask<32xpred> diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_i16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_i16.pto new file mode 100644 index 0000000000..0bd9027a50 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_i16.pto @@ -0,0 +1,36 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +module { + func.func @vmi_to_vpto_gather_i16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vreg<128xi16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xi16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xi16, #pto.vmi.layout>) + -> !pto.vreg<128xi16> + return %part : !pto.vreg<128xi16> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_i16( +// CHECK: %[[GATHER_I16:.*]] = pto.vgather2 %arg0, %arg1, %arg2 : !pto.ptr, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xi16> +// CHECK: %[[OUT_I16:.*]] = pto.vsel %[[GATHER_I16]], %{{.*}}, %arg2 : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask -> !pto.vreg<128xi16> +// CHECK: return %[[OUT_I16]] + +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16.pto new file mode 100644 index 0000000000..cf07fcf462 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16.pto @@ -0,0 +1,57 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +module { + func.func @vmi_to_vpto_gather_ui8_to_ui16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vreg<128xui16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xui16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xui16, #pto.vmi.layout>) + -> !pto.vreg<128xui16> + return %part : !pto.vreg<128xui16> + } + + func.func @vmi_to_vpto_gather_i8_to_i16( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<32xb16, #pto.vmi.layout>) + -> !pto.vreg<128xi16> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<32xui16, #pto.vmi.layout>, + !pto.vmi.mask<32xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<32xi16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<32xi16, #pto.vmi.layout>) + -> !pto.vreg<128xi16> + return %part : !pto.vreg<128xi16> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_ui8_to_ui16( +// CHECK: %[[GATHER_UI8:.*]] = pto.vgather2 %arg0, %arg1, %arg2 : !pto.ptr, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xui16> +// CHECK: %[[OUT_UI8:.*]] = pto.vsel %[[GATHER_UI8]], %{{.*}}, %arg2 : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xui16> +// CHECK: return %[[OUT_UI8]] + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_i8_to_i16( +// CHECK: %[[GATHER_I8:.*]] = pto.vgather2 %arg0, %arg1, %arg2 : !pto.ptr, !pto.vreg<128xui16>, !pto.mask -> !pto.vreg<128xi16> +// CHECK: %[[OUT_I8:.*]] = pto.vsel %[[GATHER_I8]], %{{.*}}, %arg2 : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask -> !pto.vreg<128xi16> +// CHECK: return %[[OUT_I8]] + +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16_multi_chunk.pto b/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16_multi_chunk.pto new file mode 100644 index 0000000000..5cc4b304d9 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_i8_to_i16_multi_chunk.pto @@ -0,0 +1,62 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +module { + func.func @vmi_to_vpto_gather_i8_to_i16_256( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<256xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<256xui16, #pto.vmi.layout>, + !pto.vmi.mask<256xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<256xui16, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<256xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1 : !pto.vreg<128xui16>, !pto.vreg<128xui16> + } + + func.func @vmi_to_vpto_gather_i8_to_i16_512( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<512xb16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<512xui16, #pto.vmi.layout>, + !pto.vmi.mask<512xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<512xui16, #pto.vmi.layout> + %p0, %p1, %p2, %p3 = "pto.vmi.unpack"(%out) + : (!pto.vmi.vreg<512xui16, #pto.vmi.layout>) + -> (!pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>) + return %p0, %p1, %p2, %p3 + : !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16>, !pto.vreg<128xui16> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_gather_i8_to_i16_256( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK-LABEL: func.func @vmi_to_vpto_gather_i8_to_i16_512( +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK: pto.vgather2 +// CHECK: pto.vsel +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_scatter_shape_invalid.pto b/test/lit/vmi_new/vmi_to_vpto_gather_scatter_shape_invalid.pto index ea3b408a42..3531c3b0e6 100644 --- a/test/lit/vmi_new/vmi_to_vpto_gather_scatter_shape_invalid.pto +++ b/test/lit/vmi_new/vmi_to_vpto_gather_scatter_shape_invalid.pto @@ -22,7 +22,7 @@ module { } } -// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only +// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2/pto.vgather2_bc + pto.vsel only // CHECK-SAME: contiguous result, indices, passthru, and mask layouts // ----- @@ -41,7 +41,7 @@ module { } } -// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only +// CHECK: VMI{{-}}UNSUPPORTED{{:}} pto.vmi.gather lowers through pto.vgather2/pto.vgather2_bc + pto.vsel only // CHECK-SAME: result requires full physical chunks // CHECK-SAME: found padding lane in physical chunk diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_i32.pto b/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_i32.pto new file mode 100644 index 0000000000..b50cb8058f --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_i32.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 root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto 2>&1 | FileCheck %s + +module { + func.func @gather_i32_5chunk_invalid( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<320xi32, #pto.vmi.layout>, + %mask: !pto.vmi.mask<320xb32, #pto.vmi.layout>) + -> !pto.vmi.vreg<320xi32, #pto.vmi.layout> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<320xi32, #pto.vmi.layout>, + !pto.vmi.mask<320xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<320xi32, #pto.vmi.layout> + return %out : !pto.vmi.vreg<320xi32, #pto.vmi.layout> + } +} + +// CHECK: VMI-UNSUPPORTED +// CHECK: exceeds the 4 physical register limit per VMI instruction diff --git a/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_ui16.pto b/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_ui16.pto new file mode 100644 index 0000000000..9540f38c32 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_gather_too_many_chunks_invalid_ui16.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 root of the software repository for the full text of the License. + +// RUN: not pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto 2>&1 | FileCheck %s + +module { + func.func @gather_ui16_5chunk_invalid( + %src: !pto.ptr, + %indices: !pto.vmi.vreg<640xui16, #pto.vmi.layout>, + %mask: !pto.vmi.mask<640xb16, #pto.vmi.layout>) + -> !pto.vmi.vreg<640xui16, #pto.vmi.layout> { + %out = pto.vmi.vgather %src, %indices, %mask + : !pto.ptr, + !pto.vmi.vreg<640xui16, #pto.vmi.layout>, + !pto.vmi.mask<640xb16, #pto.vmi.layout> + -> !pto.vmi.vreg<640xui16, #pto.vmi.layout> + return %out : !pto.vmi.vreg<640xui16, #pto.vmi.layout> + } +} + +// CHECK: VMI-UNSUPPORTED +// CHECK: exceeds the 4 physical register limit per VMI instruction diff --git a/test/lit/vpto/vpto_mask_simplify.pto b/test/lit/vpto/vpto_mask_simplify.pto index d238414a2a..42ea1dfd2d 100644 --- a/test/lit/vpto/vpto_mask_simplify.pto +++ b/test/lit/vpto/vpto_mask_simplify.pto @@ -83,6 +83,47 @@ module { return %low, %high, %lhs, %rhs : !pto.mask, !pto.mask, !pto.mask, !pto.mask } + func.func @vsel_all_true_b16_pset(%a: !pto.vreg<128xi16>, + %b: !pto.vreg<128xi16>) + -> !pto.vreg<128xi16> { + %mask = pto.pset_b16 "PAT_ALL" : !pto.mask + %r = pto.vsel %a, %b, %mask + : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask + -> !pto.vreg<128xi16> + return %r : !pto.vreg<128xi16> + } + + func.func @vsel_all_true_b16_pge(%a: !pto.vreg<128xi16>, + %b: !pto.vreg<128xi16>) + -> !pto.vreg<128xi16> { + %mask = pto.pge_b16 "PAT_ALL" : !pto.mask + %r = pto.vsel %a, %b, %mask + : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask + -> !pto.vreg<128xi16> + return %r : !pto.vreg<128xi16> + } + + func.func @vsel_all_true_b32_pset(%a: !pto.vreg<64xi32>, + %b: !pto.vreg<64xi32>) + -> !pto.vreg<64xi32> { + %mask = pto.pset_b32 "PAT_ALL" : !pto.mask + %r = pto.vsel %a, %b, %mask + : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask + -> !pto.vreg<64xi32> + return %r : !pto.vreg<64xi32> + } + + func.func @vsel_nonuniform_b16_preserved(%a: !pto.vreg<128xi16>, + %b: !pto.vreg<128xi16>) + -> !pto.vreg<128xi16> { + %mask = pto.pge_b16 "PAT_VL8" : !pto.mask + %r = pto.vsel %a, %b, %mask + : !pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask + -> !pto.vreg<128xi16> + return %r : !pto.vreg<128xi16> + } + + } // CHECK-LABEL: func.func @pintlv_b8_pset @@ -122,3 +163,19 @@ module { // CHECK: %[[RHS:.*]] = pto.pge_b16 "PAT_ALL" // CHECK-NOT: pto.pintlv_b16 // CHECK: return %[[LHS]], %[[RHS]], %[[LHS]], %[[RHS]] + +// CHECK-LABEL: func.func @vsel_all_true_b16_pset +// CHECK-NOT: pto.vsel +// CHECK: return %{{.*}} : !pto.vreg<128xi16> + +// CHECK-LABEL: func.func @vsel_all_true_b16_pge +// CHECK-NOT: pto.vsel +// CHECK: return %{{.*}} : !pto.vreg<128xi16> + +// CHECK-LABEL: func.func @vsel_all_true_b32_pset +// CHECK-NOT: pto.vsel +// CHECK: return %{{.*}} : !pto.vreg<64xi32> + +// CHECK-LABEL: func.func @vsel_nonuniform_b16_preserved +// CHECK: pto.vsel +// CHECK: return %{{.*}} : !pto.vreg<128xi16> diff --git a/test/vpto/cases/vmi_new/vgather-dsl-i8-ui8-256.py b/test/vpto/cases/vmi_new/vgather-dsl-i8-ui8-256.py new file mode 100644 index 0000000000..b86beb97fb --- /dev/null +++ b/test/vpto/cases/vmi_new/vgather-dsl-i8-ui8-256.py @@ -0,0 +1,166 @@ +#!/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 Python DSL cases for the 256-lane vgather 8-bit promotion paths. + +These exercise ``pto.vmi.vgather`` directly and cover the i8 -> i16 and +ui8 -> ui16 result-type inference added with the vgather board-test commit. +""" + +from pathlib import Path +import sys + +import numpy as np + + +def _bootstrap_dsl_st_common() -> None: + here = Path(__file__).resolve() + for candidate in here.parents: + common_dir = candidate / "test" / "dsl-st" + if (common_dir / "common.py").exists(): + sys.path.insert(0, str(common_dir)) + return + raise RuntimeError("Unable to locate test/dsl-st/common.py from vgather-dsl-i8-ui8-256.py") + + +_bootstrap_dsl_st_common() + +from common import auto_main, golden_output_case +from ptodsl import pto + + +ELEMS = 256 +SRC_BYTES = 256 +IDX_BYTES = 512 +DST_BYTES = 512 + + +@pto.jit( + name="vmi_vgather_i8_to_i16_256_kernel", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def vmi_vgather_i8_to_i16_256_kernel( + src_gm: pto.ptr(pto.i8, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.i16, "gm"), +): + ub_src = pto.castptr(pto.const(0, dtype=pto.i64), pto.ptr(pto.i8, "ub")) + ub_idx = pto.castptr(pto.const(4096, dtype=pto.i64), pto.ptr(pto.ui16, "ub")) + ub_dst = pto.castptr(pto.const(8192, dtype=pto.i64), pto.ptr(pto.i16, "ub")) + + pto.mte_gm_ub(src_gm, ub_src, 0, SRC_BYTES, nburst=(1, SRC_BYTES, SRC_BYTES)) + pto.mte_gm_ub(idx_gm, ub_idx, 0, IDX_BYTES, nburst=(1, IDX_BYTES, IDX_BYTES)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + offset = pto.const(0, dtype=pto.index) + idx = pto.vmi.vload(ub_idx, offset, size=ELEMS) + mask = pto.vmi.create_mask(pto.const(ELEMS, dtype=pto.index), size=ELEMS) + out = pto.vmi.vgather(ub_src, idx, mask) + pto.vmi.vstore(out, ub_dst, offset, mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(ub_dst, dst_gm, DST_BYTES, nburst=(1, DST_BYTES, DST_BYTES)) + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.jit( + name="vmi_vgather_u8_to_u16_256_kernel", + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, +) +def vmi_vgather_u8_to_u16_256_kernel( + src_gm: pto.ptr(pto.ui8, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.ui16, "gm"), +): + ub_src = pto.castptr(pto.const(0, dtype=pto.i64), pto.ptr(pto.ui8, "ub")) + ub_idx = pto.castptr(pto.const(4096, dtype=pto.i64), pto.ptr(pto.ui16, "ub")) + ub_dst = pto.castptr(pto.const(8192, dtype=pto.i64), pto.ptr(pto.ui16, "ub")) + + pto.mte_gm_ub(src_gm, ub_src, 0, SRC_BYTES, nburst=(1, SRC_BYTES, SRC_BYTES)) + pto.mte_gm_ub(idx_gm, ub_idx, 0, IDX_BYTES, nburst=(1, IDX_BYTES, IDX_BYTES)) + + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + offset = pto.const(0, dtype=pto.index) + idx = pto.vmi.vload(ub_idx, offset, size=ELEMS) + mask = pto.vmi.create_mask(pto.const(ELEMS, dtype=pto.index), size=ELEMS) + out = pto.vmi.vgather(ub_src, idx, mask) + pto.vmi.vstore(out, ub_dst, offset, mask) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=0) + pto.mte_ub_gm(ub_dst, dst_gm, DST_BYTES, nburst=(1, DST_BYTES, DST_BYTES)) + pto.pipe_barrier(pto.Pipe.ALL) + + +def _non_identity_indices() -> np.ndarray: + return ((np.arange(ELEMS, dtype=np.uint16) * 7 + 13) % ELEMS).astype(np.uint16) + + +def _i8_inputs(): + src = np.arange(ELEMS, dtype=np.uint8) + return [src, _non_identity_indices()] + + +def _i8_expected(src, idx): + return src.astype(np.uint16)[idx].astype(np.uint16) + + +def _u8_inputs(): + src = np.arange(ELEMS, dtype=np.uint8) + return [src, _non_identity_indices()] + + +def _u8_expected(src, idx): + return src.astype(np.uint16)[idx].astype(np.uint16) + + +CASES = [ + golden_output_case( + "vmi_vgather_i8_to_i16_256", + vmi_vgather_i8_to_i16_256_kernel, + inputs=_i8_inputs, + expected=_i8_expected, + output_shape=(ELEMS,), + output_dtype=np.uint16, + rtol=0.0, + atol=0.0, + ), + golden_output_case( + "vmi_vgather_u8_to_u16_256", + vmi_vgather_u8_to_u16_256_kernel, + inputs=_u8_inputs, + expected=_u8_expected, + output_shape=(ELEMS,), + output_dtype=np.uint16, + rtol=0.0, + atol=0.0, + ), +] + + +KERNELS = [ + vmi_vgather_i8_to_i16_256_kernel, + vmi_vgather_u8_to_u16_256_kernel, +] + + +auto_main(globals()) diff --git a/test/vpto/cases/vmi_new/vgather-vmi.py b/test/vpto/cases/vmi_new/vgather-vmi.py new file mode 100644 index 0000000000..52fa6ce942 --- /dev/null +++ b/test/vpto/cases/vmi_new/vgather-vmi.py @@ -0,0 +1,516 @@ +#!/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 source-backed cases for the 512-lane vgather VMI kernels. + +The MLIR is embedded inline so these cases are self-contained and do not depend +on the old board-test directories. +""" + +from pathlib import Path +import sys + +import numpy as np + + +def _bootstrap_dsl_st_common() -> None: + here = Path(__file__).resolve() + for candidate in here.parents: + common_dir = candidate / "test" / "dsl-st" + if (common_dir / "common.py").exists(): + sys.path.insert(0, str(common_dir)) + return + raise RuntimeError("Unable to locate test/dsl-st/common.py from vgather-vmi.py") + + +_bootstrap_dsl_st_common() + +from common import auto_main, golden_output_case +from ptodsl import pto + + +BF16_SRC = r'''// 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. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_vgather_bf16_to_bf16_512_kernel(%src_gm: !pto.ptr, + %idx_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_idx = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %idx_gm, %ub_idx, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %idx = pto.vmi.vload %ub_idx[%c0] + : !pto.ptr -> !pto.vmi.vreg<512xui16> + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xpred> + %out = pto.vmi.vgather %ub_src, %idx, %mask + : !pto.ptr, !pto.vmi.vreg<512xui16>, + !pto.vmi.mask<512xpred> -> !pto.vmi.vreg<512xbf16> + pto.vmi.vstore %out, %ub_dst[%c0] + : !pto.vmi.vreg<512xbf16>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} +''' +F16_SRC = r'''// 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. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_vgather_f16_to_f16_512_kernel(%src_gm: !pto.ptr, + %idx_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_idx = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %idx_gm, %ub_idx, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %idx = pto.vmi.vload %ub_idx[%c0] + : !pto.ptr -> !pto.vmi.vreg<512xui16> + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xpred> + %out = pto.vmi.vgather %ub_src, %idx, %mask + : !pto.ptr, !pto.vmi.vreg<512xui16>, + !pto.vmi.mask<512xpred> -> !pto.vmi.vreg<512xf16> + pto.vmi.vstore %out, %ub_dst[%c0] + : !pto.vmi.vreg<512xf16>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} +''' +I8_512_SRC = r'''// 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. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_vgather_i8_to_i16_512_kernel(%src_gm: !pto.ptr, + %idx_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_idx = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c512_i64 + nburst(%c1_i64, %c512_i64, %c512_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %idx_gm, %ub_idx, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %idx = pto.vmi.vload %ub_idx[%c0] + : !pto.ptr -> !pto.vmi.vreg<512xui16> + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xpred> + %out = pto.vmi.vgather %ub_src, %idx, %mask + : !pto.ptr, !pto.vmi.vreg<512xui16>, + !pto.vmi.mask<512xpred> -> !pto.vmi.vreg<512xi16> + pto.vmi.vstore %out, %ub_dst[%c0] + : !pto.vmi.vreg<512xi16>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} +''' +U8_512_SRC = r'''// 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. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_vgather_u8_to_u16_512_kernel(%src_gm: !pto.ptr, + %idx_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_idx = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c512_i64 + nburst(%c1_i64, %c512_i64, %c512_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %idx_gm, %ub_idx, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %idx = pto.vmi.vload %ub_idx[%c0] + : !pto.ptr -> !pto.vmi.vreg<512xui16> + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xpred> + %out = pto.vmi.vgather %ub_src, %idx, %mask + : !pto.ptr, !pto.vmi.vreg<512xui16>, + !pto.vmi.mask<512xpred> -> !pto.vmi.vreg<512xui16> + pto.vmi.vstore %out, %ub_dst[%c0] + : !pto.vmi.vreg<512xui16>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} +''' +UI16_512_SRC = r'''// 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. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @vmi_vgather_ui16_to_ui16_512_kernel(%src_gm: !pto.ptr, + %idx_gm: !pto.ptr, + %dst_gm: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c512_i64 = arith.constant 512 : i64 + %c1024_i64 = arith.constant 1024 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %c8192_i64 = arith.constant 8192 : i64 + + %ub_src = pto.castptr %c0_i64 : i64 -> !pto.ptr + %ub_idx = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %ub_dst = pto.castptr %c8192_i64 : i64 -> !pto.ptr + + pto.mte_gm_ub %src_gm, %ub_src, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + pto.mte_gm_ub %idx_gm, %ub_idx, %c0_i64, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + + pto.set_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_V", "EVENT_ID0"] + + pto.vecscope { + %idx = pto.vmi.vload %ub_idx[%c0] + : !pto.ptr -> !pto.vmi.vreg<512xui16> + %mask = pto.vmi.create_mask %c512 : index -> !pto.vmi.mask<512xpred> + %out = pto.vmi.vgather %ub_src, %idx, %mask + : !pto.ptr, !pto.vmi.vreg<512xui16>, + !pto.vmi.mask<512xpred> -> !pto.vmi.vreg<512xui16> + pto.vmi.vstore %out, %ub_dst[%c0] + : !pto.vmi.vreg<512xui16>, !pto.ptr + } + + pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] + pto.mte_ub_gm %ub_dst, %dst_gm, %c1024_i64 + nburst(%c1_i64, %c1024_i64, %c1024_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + pto.barrier #pto.pipe + return + } +} +''' +@pto.jit( + name="vmi_vgather_bf16_to_bf16_512_kernel", + target="a5", + backend="vpto", + mode="explicit", + source=BF16_SRC, +) +def vmi_vgather_bf16_to_bf16_512_kernel( + src_gm: pto.ptr(pto.bf16, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.bf16, "gm"), +): + pass + + +@pto.jit( + name="vmi_vgather_f16_to_f16_512_kernel", + target="a5", + backend="vpto", + mode="explicit", + source=F16_SRC, +) +def vmi_vgather_f16_to_f16_512_kernel( + src_gm: pto.ptr(pto.f16, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.f16, "gm"), +): + pass + + +@pto.jit( + name="vmi_vgather_i8_to_i16_512_kernel", + target="a5", + backend="vpto", + mode="explicit", + source=I8_512_SRC, +) +def vmi_vgather_i8_to_i16_512_kernel( + src_gm: pto.ptr(pto.i8, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.i16, "gm"), +): + pass + + +@pto.jit( + name="vmi_vgather_u8_to_u16_512_kernel", + target="a5", + backend="vpto", + mode="explicit", + source=U8_512_SRC, +) +def vmi_vgather_u8_to_u16_512_kernel( + src_gm: pto.ptr(pto.ui8, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.ui16, "gm"), +): + pass + + +@pto.jit( + name="vmi_vgather_ui16_to_ui16_512_kernel", + target="a5", + backend="vpto", + mode="explicit", + source=UI16_512_SRC, +) +def vmi_vgather_ui16_to_ui16_512_kernel( + src_gm: pto.ptr(pto.ui16, "gm"), + idx_gm: pto.ptr(pto.ui16, "gm"), + dst_gm: pto.ptr(pto.ui16, "gm"), +): + pass + + +def _non_identity_indices(elems: int) -> np.ndarray: + # Non-identity deterministic gather (idx[i] != i for most lanes). + return ((np.arange(elems, dtype=np.uint16) * 7 + 13) % elems).astype(np.uint16) + + +def _bf16_512_inputs(): + elems = 512 + src = np.arange(elems, dtype=np.float16) + return [src, _non_identity_indices(elems)] + + +def _bf16_512_expected(src, idx): + return src[idx].astype(np.float16) + + +def _f16_512_inputs(): + elems = 512 + src = np.arange(elems, dtype=np.float16) + return [src, _non_identity_indices(elems)] + + +def _f16_512_expected(src, idx): + return src[idx].astype(np.float16) + + +def _i8_512_inputs(): + elems = 512 + src = np.arange(elems, dtype=np.uint8) + return [src, _non_identity_indices(elems)] + + +def _i8_512_expected(src, idx): + return src.astype(np.uint16)[idx].astype(np.uint16) + + +def _u8_512_inputs(): + elems = 512 + src = np.arange(elems, dtype=np.uint8) + return [src, _non_identity_indices(elems)] + + +def _u8_512_expected(src, idx): + return src.astype(np.uint16)[idx].astype(np.uint16) + + +def _ui16_512_inputs(): + elems = 512 + src = np.arange(elems, dtype=np.uint16) + return [src, _non_identity_indices(elems)] + + +def _ui16_512_expected(src, idx): + return src[idx].astype(np.uint16) + + +CASES = [ + golden_output_case( + "vmi_vgather_bf16_to_bf16_512", + vmi_vgather_bf16_to_bf16_512_kernel, + inputs=_bf16_512_inputs, + expected=_bf16_512_expected, + output_shape=(512,), + output_dtype=np.float16, + rtol=0.0, + atol=0.0, + ), + golden_output_case( + "vmi_vgather_f16_to_f16_512", + vmi_vgather_f16_to_f16_512_kernel, + inputs=_f16_512_inputs, + expected=_f16_512_expected, + output_shape=(512,), + output_dtype=np.float16, + rtol=0.0, + atol=0.0, + ), + golden_output_case( + "vmi_vgather_i8_to_i16_512", + vmi_vgather_i8_to_i16_512_kernel, + inputs=_i8_512_inputs, + expected=_i8_512_expected, + output_shape=(512,), + output_dtype=np.uint16, + rtol=0.0, + atol=0.0, + ), + golden_output_case( + "vmi_vgather_u8_to_u16_512", + vmi_vgather_u8_to_u16_512_kernel, + inputs=_u8_512_inputs, + expected=_u8_512_expected, + output_shape=(512,), + output_dtype=np.uint16, + rtol=0.0, + atol=0.0, + ), + golden_output_case( + "vmi_vgather_ui16_to_ui16_512", + vmi_vgather_ui16_to_ui16_512_kernel, + inputs=_ui16_512_inputs, + expected=_ui16_512_expected, + output_shape=(512,), + output_dtype=np.uint16, + rtol=0.0, + atol=0.0, + ), +] + + +KERNELS = [ + vmi_vgather_bf16_to_bf16_512_kernel, + vmi_vgather_f16_to_f16_512_kernel, + vmi_vgather_i8_to_i16_512_kernel, + vmi_vgather_u8_to_u16_512_kernel, + vmi_vgather_ui16_to_ui16_512_kernel, +] + + +auto_main(globals())