diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 75873dcb7a..eca1494909 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5458,6 +5458,51 @@ FailureOr createIotaContiguousChunk(Location loc, Type resultType, .getResult(); } +// Materialize a contiguous logical iota into a lane-strided physical chunk. +// On the A5 VPTO path VCI over a narrow floating-point carrier advances in +// physical lanes. The lane-stride store then selects every `laneStride`-th +// lane, so the VCI ramp has to be scaled by the reciprocal stride to preserve +// a unit logical increment. +FailureOr createIotaLaneStrideChunk( + Location loc, Type resultType, Value base, int64_t laneStride, int64_t laneOffset, StringAttr orderAttr, + PatternRewriter& rewriter) +{ + auto vregType = dyn_cast(resultType); + if (!vregType || (laneStride != 2 && laneStride != 4)) { + return failure(); + } + FailureOr mask = createAllTrueMaskForVReg(loc, vregType, rewriter); + FailureOr zero = createScalarOffsetConstant(loc, base.getType(), 0, rewriter); + // This materialization is needed for narrow floating-point carriers (the + // f16 PK_B32 case in particular). Integer VCI uses its native unit-lane + // semantics and must not be scaled here. + int64_t scale = 1; + if (isa(base.getType())) + scale = laneStride; + FailureOr factor; + if (isa(base.getType())) { + auto floatType = cast(base.getType()); + factor = rewriter.create(loc, rewriter.getFloatAttr(floatType, 1.0 / scale)).getResult(); + } else { + factor = createScalarOffsetConstant(loc, base.getType(), 1, rewriter); + } + if (failed(mask) || failed(zero) || failed(factor)) { + return failure(); + } + Value indices = rewriter.create(loc, resultType, *zero, StringAttr{}).getResult(); + Value scaled = rewriter.create(loc, resultType, indices, *factor, *mask).getResult(); + StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); + FailureOr chunkBase = createIotaChunkBase(loc, base, laneOffset, order, rewriter); + if (failed(chunkBase)) { + return failure(); + } + if (order == "DESC") { + Value baseVector = rewriter.create(loc, resultType, *chunkBase, *mask, nullptr).getResult(); + return rewriter.create(loc, resultType, baseVector, scaled, *mask).getResult(); + } + return rewriter.create(loc, resultType, scaled, *chunkBase, *mask).getResult(); +} + /// Pack group-periodic ramps inside one physical VL when S < physVL and /// physVL % S == 0 (e.g. i32 L=64,group=2 → [base..base+31 | base..base+31]). /// @@ -5751,6 +5796,28 @@ struct OneToNVMIIotaOpPattern : OneToNOpConversionPattern { return success(); } + if (layout.isContiguous() && layout.getLaneStride() != 1) { + int64_t laneStride = layout.getLaneStride(); + if (laneStride != 2 && laneStride != 4) { + return rewriter.notifyMatchFailure(op, "unsupported contiguous iota lane_stride"); + } + int64_t logicalLanesPerChunk = *lanesPerPart / laneStride; + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) { + return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + } + FailureOr result = createIotaLaneStrideChunk( + op.getLoc(), resultType, *base, laneStride, static_cast(index) * logicalLanesPerChunk, + op.getOrderAttr(), rewriter); + if (failed(result)) { + return rewriter.notifyMatchFailure(op, "failed to materialize lane-strided iota chunk"); + } + results.push_back(*result); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); + return success(); + } + if (layout.isContiguous()) { for (auto [index, resultType] : llvm::enumerate(resultTypes)) { if (!isa(resultType)) diff --git a/packaging/ptoas-vmi/pyproject.toml.patch b/packaging/ptoas-vmi/pyproject.toml.patch index f2f71ec41c..87dc5d8c72 100644 --- a/packaging/ptoas-vmi/pyproject.toml.patch +++ b/packaging/ptoas-vmi/pyproject.toml.patch @@ -10,7 +10,7 @@ diff --git a/pyproject.toml b/pyproject.toml +version = "0.1.6" +description = "PTO Assembler & Optimizer with VMI support" readme = "README.md" - requires-python = ">=3.10" + requires-python = ">=3.7" license = "Apache-2.0" @@ -43,4 +43,5 @@ cmake.version = "CMakeLists.txt" ninja.version = ">=1.10" diff --git a/test/lit/vmi_new/vmi_to_vpto_iota.pto b/test/lit/vmi_new/vmi_to_vpto_iota.pto index 79f8443c48..8652c76170 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota.pto @@ -49,6 +49,19 @@ module { return %part : !pto.vreg<128xi16> } + // Regression for #1411: a compact f16 store chooses a lane_stride=2 + // carrier. The physical ramp must advance by two so PK_B32 observes a + // logical step of one. + func.func @vmi_to_vpto_iota_f16_lane_stride2(%base: f16) + -> !pto.vreg<128xf16> { + %value = pto.vmi.vci %base + : f16 -> !pto.vmi.vreg<64xf16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<64xf16, #pto.vmi.layout>) + -> !pto.vreg<128xf16> + return %part : !pto.vreg<128xf16> + } + func.func @vmi_to_vpto_iota_f16_deint2_asc(%base: f16) -> (!pto.vreg<128xf16>, !pto.vreg<128xf16>) { %value = pto.vmi.vci %base @@ -101,6 +114,15 @@ module { // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast +// CHECK-LABEL: func.func @vmi_to_vpto_iota_f16_lane_stride2( +// CHECK: %[[C2:.*]] = arith.constant 5.000000e-01 : f16 +// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f16 +// CHECK: %[[IDX:.*]] = pto.vci %[[C0]] : f16 -> !pto.vreg<128xf16> +// CHECK: %[[SCALED:.*]] = pto.vmuls %[[IDX]], %[[C2]] +// CHECK: %[[RESULT:.*]] = pto.vadds %[[SCALED]], %arg0 +// CHECK: return %[[RESULT]] +// CHECK-NOT: pto.vmi. + // CHECK-LABEL: func.func @vmi_to_vpto_iota_f16_deint2_asc( // CHECK: arith.constant 1.000000e+00 : f16 // CHECK: arith.constant 2.000000e+00 : f16