diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 186e0b7eff..e9add2704c 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -412,7 +412,32 @@ struct LayoutSolver { if (failed(fact)) { return {}; } - return fact->layout.resultLayout; + + VMILayoutAttr directLayout = fact->layout.resultLayout; + // A direct E2B packet fills exactly one physical part. When the + // contiguous form of this broadcast spans multiple physical chunks, the + // direct table can only offer a deinterleaved split layout. Prefer the + // generic contiguous lowering (group_slots -> contiguous) so consumers + // such as plain elementwise vmul can stay contiguous and avoid + // vldsx2/vintlv-style deinterleave materialization. + if (fact->kind == VMIGroupBroadcastLoadDirectKind::E2B && + directLayout.isDeinterleaved()) { + VMILayoutAttr contiguous = getContiguousLayout(); + auto contiguousType = VMIVRegType::get( + ctx, type.getElementCount(), type.getElementType(), contiguous); + FailureOr contiguousArity = getVMIPhysicalArity(contiguousType); + if (succeeded(contiguousArity) && *contiguousArity > 1) { + FailureOr< + SmallVector> + contiguousFacts = supports.getGroupBroadcastLayoutFactsForLayout( + type, contiguousType, op.getNumGroupsAttr().getInt(), + VMIGroupBroadcastLayoutPort::Result, contiguous); + if (succeeded(contiguousFacts) && !contiguousFacts->empty()) { + return contiguous; + } + } + } + return directLayout; } VMILayoutAttr getPreferredGroupBroadcastSourceLayout(Value value, @@ -928,6 +953,42 @@ struct LayoutSolver { .Case([this, op](auto load) { return addDeinterleaveLoadConstraint(load, op); }) + .Case([this, op](auto load) { + auto type = cast(load.getResult().getType()); + FailureOr lanesPerPart = + getDataLanesPerPart(type.getElementType()); + // Stock leaves plain VMILoadOp results unconstrained (the layout is + // driven by consumer use-requests). Seed a natural layout ONLY + // when the packed/small-load lane-stride heuristic below actually + // applies: seeding plain contiguous unconditionally would block + // deinterleaved use-requests from propagating back to the load + // (e.g. channel_split @128 f16 hitting the ensure-layout gap, + // vmi_layout_assignment_store_prefer_lane_stride). + // The heuristic is safe for packed float carriers (f4x2 / + // hi-float8x2 / bf16x2, whose storage byte holds a pair of values) + // and for >=16-bit dense elements (a lane-stride unpack keeps every + // real element reachable through the EVEN part; pinned by the + // ComputeMropeF16 capability guard). Dense 8-bit elements are the + // proven mine: an unpack-style distribution whose gap lanes read as + // zeros (the e4m3 odd-column regression shape; hidden=128 configs + // would hit it), so those keep no seed at all. + Type loadElemTy = type.getElementType(); + bool isPackedCarrier = pto::isPTOFloat4PackedType(loadElemTy) || + pto::isPTOHiFloat8x2Type(loadElemTy) || + pto::isPTOBF16x2Type(loadElemTy); + bool laneStrideSeedAllowed = + isPackedCarrier || + pto::getPTOStorageElemBitWidth(loadElemTy) != 8; + if (!laneStrideSeedAllowed || failed(lanesPerPart) || + type.getElementCount() >= *lanesPerPart || + *lanesPerPart % type.getElementCount() != 0) { + return std::optional(std::nullopt); + } + int64_t laneStride = *lanesPerPart / type.getElementCount(); + return constraintResult(setNaturalLayout( + load.getResult(), VMILayoutAttr::getContiguous(ctx, laneStride), + op)); + }) .Case([this, op](auto load) { requestDataUse(load.getPassthruMutable(), getContiguousLayout()); return constraintResult(setNaturalLayout( diff --git a/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc b/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc index bc23a7964b..b10fdc1f51 100644 --- a/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc +++ b/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc @@ -28,6 +28,11 @@ enum class CastTypeClass { Float, // Both source and result are integer element types. Integer, + // At least one side is a packed float carrier (f4x2 / hi-float8x2 / + // bf16x2): a storage byte that holds a pair of values. Lane-stride + // forms tuned for paired data (UNPK-style distributions) must not be + // selected for dense element casts (see the e4m3 odd-column regression). + PackedFloatCarrier, }; struct LayoutPattern { @@ -138,6 +143,13 @@ static bool matchesCastTypeClass(CastTypeClass typeClass, Type sourceType, pto::isPTOLowPrecisionType(resultType)); case CastTypeClass::Integer: return isa(sourceType) && isa(resultType); + case CastTypeClass::PackedFloatCarrier: + return pto::isPTOFloat4PackedType(sourceType) || + pto::isPTOHiFloat8x2Type(sourceType) || + pto::isPTOBF16x2Type(sourceType) || + pto::isPTOFloat4PackedType(resultType) || + pto::isPTOHiFloat8x2Type(resultType) || + pto::isPTOBF16x2Type(resultType); } llvm_unreachable("unknown cast type class"); } diff --git a/lib/PTO/Transforms/VMILayoutSupportTables.inc b/lib/PTO/Transforms/VMILayoutSupportTables.inc index e74025b750..1bfb808a6a 100644 --- a/lib/PTO/Transforms/VMILayoutSupportTables.inc +++ b/lib/PTO/Transforms/VMILayoutSupportTables.inc @@ -183,6 +183,15 @@ static constexpr PreferredCastLayoutPattern kPreferredCastLayoutPatterns[] = { // compact lane-stride form is the natural cast layout. {bits<16>(), bits<32>(), 64, ls(2), c()}, {bits<8>(), bits<32>(), 64, ls(4), c()}, + // fp4 pair-cast row: f4E2M1x2(8-bit storage) -> bf16x2(32-bit pair) at 128 + // elements (UNPK4 + P0; dense-safe -- LaneExpand4 carries every byte and + // the sourceBits==8 ls(4) early path selects exactly the valid lanes). + // Restricted to packed float carriers so dense 8-bit -> 32-bit casts keep + // the stock default row {c(), d(4)}: with this row matching dense casts + // the solver preferred c() over d(4) (both value-correct, but it changes + // downstream solver outcomes and conflict diagnostics -- see the + // multi_return_conflict_invalid lit test). + {bits<8>(), bits<32>(), 128, ls(4), c(), CastTypeClass::PackedFloatCarrier}, {bits<32>(), bits<16>(), 64, c(), ls(2)}, {bits<32>(), bits<8>(), 64, c(), ls(4)}, {bits<32>(), bits<8>(), 128, d(2), ls(2)}, @@ -238,7 +247,14 @@ static constexpr LegalCastLayoutPattern kLegalCastLayoutPatterns[] = { // 4x widening/narrowing. {bits<8>(), bits<32>(), c(), d(4)}, - {bits<8>(), bits<32>(), ls(2), d(2)}, + // Integer-only: both observed Float consumers of this row are broken. + // - dense f8 (e4m3 -> f32/bf16): result d(2) back-propagates ls(2) onto + // the source; ls(2) lowers to UNPK_B8 whose zero-fill gaps surface as + // the "odd columns = 0" 50% regression. + // - packed f4x2: the original pre-578dec5d 47% bug (same shape). + // Restricting the row to Integer forces the Float solver back to the + // consistent {c(), d(4)} solution (stock form). + {bits<8>(), bits<32>(), ls(2), d(2), CastTypeClass::Integer}, {bits<8>(), bits<32>(), ls(4), c()}, {bits<32>(), bits<8>(), d(4), c()}, {bits<32>(), bits<8>(), c(), ls(4)}, diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 75873dcb7a..9c203ff9c6 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -8926,6 +8926,22 @@ struct OneToNVMIGroupBroadcastLoadOpPattern firstType, VPTOMemoryOpFamily::Load, e2bDist); } + // One E2B packet materializes exactly one physical part. A contiguous + // broadcast result that spans multiple physical chunks has no single + // reusable packet per part, so keep the direct-E2B path only for the + // one-packet-per-part shapes and let the generic group_slots -> + // contiguous broadcast fallback handle the rest. + if (canUseDirectE2B) { + VMILayoutAttr resultLayout = resultVMIType.getLayoutAttr(); + FailureOr contiguousChunksPerPart = + getDataChunksInPart(resultVMIType, 0); + if (resultLayout && resultLayout.isContiguous() && + (failed(contiguousChunksPerPart) || + *contiguousChunksPerPart != 1)) { + canUseDirectE2B = false; + } + } + if (failed(directFact) || directFact->kind != VMIGroupBroadcastLoadDirectKind::E2B || !canUseDirectE2B) { @@ -11306,15 +11322,42 @@ struct OneToNVMIExtFOpPattern : OneToNOpConversionPattern { ArrayRef parts; int64_t factor = 0; - if (sourceBits == 16 && resultTypes.size() == 2 * sourceParts.size()) { + if (sourceBits == 16) { static constexpr StringRef kEvenOddParts[] = {"EVEN", "ODD"}; - parts = kEvenOddParts; - factor = 2; - } else if (sourceBits == 8 && - resultTypes.size() == 4 * sourceParts.size()) { + constexpr int64_t kMaxFactor = 2; + if (resultTypes.size() == 0 || + resultTypes.size() % sourceParts.size() != 0 || + resultTypes.size() > kMaxFactor * sourceParts.size()) { + return rewriter.notifyMatchFailure( + op, "unsupported physical extf source/result width relation"); + } + factor = resultTypes.size() / sourceParts.size(); + parts = ArrayRef(kEvenOddParts, factor); + } else if (sourceBits == 8) { static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"}; - parts = kPacked4Parts; - factor = 4; + constexpr int64_t kMaxFactor = 4; + if (resultTypes.size() == 0 || + resultTypes.size() % sourceParts.size() != 0 || + resultTypes.size() > kMaxFactor * sourceParts.size()) { + return rewriter.notifyMatchFailure( + op, "unsupported physical extf source/result width relation"); + } + factor = resultTypes.size() / sourceParts.size(); + // Dense 8-bit sources (f8*/i8/ui8) keep the stock factor-4 requirement: + // a sub-byte part selection (factor < 4) is only meaningful when the + // storage byte holds a pair of packed values (f4x2 / hi-float8x2 / + // bf16x2). For dense bytes, factor<4 parts would read the zero-fill + // gaps of an unpack-style distribution and silently produce wrong + // output (the e4m3 odd-column regression). Restore the pre-578dec5d + // compile-time rejection for that case. + if (!isVMIPackedFloatCarrierType(sourceVMIType.getElementType()) && + factor != kMaxFactor) { + return rewriter.notifyMatchFailure( + op, + "dense 8-bit extf requires factor 4 (one result part per source " + "byte)"); + } + parts = ArrayRef(kPacked4Parts, factor); } else { return rewriter.notifyMatchFailure( op, "unsupported physical extf source/result width relation"); @@ -11327,15 +11370,15 @@ struct OneToNVMIExtFOpPattern : OneToNOpConversionPattern { SmallVector results; results.reserve(resultTypes.size()); - for (int64_t partIndex = 0; partIndex < factor; ++partIndex) { - for (auto [chunkIndex, sourcePart] : llvm::enumerate(sourceParts)) { - VRegType resultType = - resultVRegTypes[partIndex * sourceParts.size() + chunkIndex]; - results.push_back(viewVcvtResult( - resultType, sourcePart, *mask, /*rnd=*/nullptr, /*sat=*/nullptr, - rewriter.getStringAttr(parts[partIndex]))); + for (int64_t partIndex = 0; partIndex < factor; ++partIndex) { + for (auto [chunkIndex, sourcePart] : llvm::enumerate(sourceParts)) { + VRegType resultType = + resultVRegTypes[partIndex * sourceParts.size() + chunkIndex]; + results.push_back(viewVcvtResult( + resultType, sourcePart, *mask, /*rnd=*/nullptr, /*sat=*/nullptr, + rewriter.getStringAttr(parts[partIndex]))); + } } - } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); return success(); diff --git a/test/lit/vmi_new/opt/compute_y1_to_fp8_fp16_vmi_opt.pto b/test/lit/vmi_new/opt/compute_y1_to_fp8_fp16_vmi_opt.pto index bdf3de9904..efbc717ff1 100644 --- a/test/lit/vmi_new/opt/compute_y1_to_fp8_fp16_vmi_opt.pto +++ b/test/lit/vmi_new/opt/compute_y1_to_fp8_fp16_vmi_opt.pto @@ -108,33 +108,43 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> !pto.vreg<128xf16> +// CHECK-NOT: E2B_B16 +// CHECK: pto.vsldb +// CHECK: pto.vselr // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: scf.for -// CHECK: pto.vldsx2 {{.*}} "DINTLV_B16" : !pto.ptr, index -> !pto.vreg<128xf16>, !pto.vreg<128xf16> +// CHECK: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<128xf16> +// CHECK: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<128xf16> // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: pto.vcvt {{.*}} {part = "P0", rnd = "R", sat = "SAT"} -// CHECK: pto.vcvt {{.*}} {part = "P1", rnd = "R", sat = "SAT"} // CHECK: pto.vcvt {{.*}} {part = "P2", rnd = "R", sat = "SAT"} -// CHECK: pto.vcvt {{.*}} {part = "P3", rnd = "R", sat = "SAT"} -// CHECK: pto.vor // CHECK: pto.vor +// CHECK: pto.vcvt {{.*}} {part = "P0", rnd = "R", sat = "SAT"} +// CHECK: pto.vcvt {{.*}} {part = "P2", rnd = "R", sat = "SAT"} // CHECK: pto.vor -// CHECK: pto.vsts {{.*}} : !pto.vreg<256xf8E4M3FN>, !pto.ptr, !pto.mask +// CHECK: pto.vsts {{.*}} {dist = "PK_B16"} : !pto.vreg<256xf8E4M3FN>, !pto.ptr, !pto.mask +// CHECK-NOT: E2B_B16 // CHECK-LABEL: func.func @ComputeY1ToFP8_fp16_e5m2_vmi -// CHECK: pto.vlds {{.*}} {dist = "E2B_B16"} : !pto.ptr -> !pto.vreg<128xf16> +// CHECK-NOT: E2B_B16 +// CHECK: pto.vsldb +// CHECK: pto.vselr // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: scf.for -// CHECK: pto.vldsx2 {{.*}} "DINTLV_B16" : !pto.ptr, index -> !pto.vreg<128xf16>, !pto.vreg<128xf16> +// CHECK: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<128xf16> +// CHECK: pto.vlds {{.*}} : !pto.ptr -> !pto.vreg<128xf16> +// CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<64xf32> // CHECK: pto.vcvt {{.*}} {part = "P0", rnd = "R", sat = "SAT"} -// CHECK: pto.vcvt {{.*}} {part = "P1", rnd = "R", sat = "SAT"} // CHECK: pto.vcvt {{.*}} {part = "P2", rnd = "R", sat = "SAT"} -// CHECK: pto.vcvt {{.*}} {part = "P3", rnd = "R", sat = "SAT"} -// CHECK: pto.vor // CHECK: pto.vor +// CHECK: pto.vcvt {{.*}} {part = "P0", rnd = "R", sat = "SAT"} +// CHECK: pto.vcvt {{.*}} {part = "P2", rnd = "R", sat = "SAT"} // CHECK: pto.vor -// CHECK: pto.vsts {{.*}} : !pto.vreg<256xf8E5M2>, !pto.ptr, !pto.mask +// CHECK: pto.vsts {{.*}} {dist = "PK_B16"} : !pto.vreg<256xf8E5M2>, !pto.ptr, !pto.mask +// CHECK-NOT: E2B_B16 diff --git a/test/lit/vmi_new/vmi_layout_assignment_group_slot_broadcast_load_e2b_b16.pto b/test/lit/vmi_new/vmi_layout_assignment_group_slot_broadcast_load_e2b_b16.pto index a08f1bcee0..c0c1b0219e 100644 --- a/test/lit/vmi_new/vmi_layout_assignment_group_slot_broadcast_load_e2b_b16.pto +++ b/test/lit/vmi_new/vmi_layout_assignment_group_slot_broadcast_load_e2b_b16.pto @@ -7,6 +7,12 @@ // 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-pre-assignment-combine -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s +// +// Multi-part group-broadcast loads prefer the contiguous fallback +// (vsldb + vselr materialization) since the H1 seed change; direct E2B_B16/B32 +// loads remain available for single-part broadcasts (see the e2b_b32 function +// below) and for operands pinned to an explicit deinterleaved layout (see +// vmi_to_vpto_group_broadcast_load_e2b_b16.pto). module { func.func @vmi_layout_assignment_group_slot_broadcast_load_e2b_b16( @@ -67,8 +73,11 @@ module { // CHECK-LABEL: func.func @vmi_layout_assignment_group_slot_broadcast_load_e2b_b16_deint2 // CHECK-SAME: (%[[SRC2:.*]]: !pto.ptr, %[[OFF2:.*]]: index) -// CHECK: %[[E2B2:.*]] = pto.vlds %[[SRC2]][{{.*}}] {dist = "E2B_B16"} : !pto.ptr -> !pto.vreg<128xbf16> -// CHECK: return %[[E2B2]], %[[E2B2]] : !pto.vreg<128xbf16>, !pto.vreg<128xbf16> +// CHECK-NOT: E2B_B16 +// CHECK: pto.vsldb +// CHECK: pto.vselr +// CHECK: return {{%.*}}, {{%.*}} : !pto.vreg<128xbf16>, !pto.vreg<128xbf16> +// CHECK-NOT: E2B_B16 // CHECK-LABEL: func.func @vmi_layout_assignment_group_slot_broadcast_load_e2b_b32 // CHECK-SAME: (%[[SRC3:.*]]: !pto.ptr, %[[OFF3:.*]]: index) @@ -77,5 +86,8 @@ module { // CHECK-LABEL: func.func @vmi_layout_assignment_group_slot_broadcast_load_e2b_b32_deint2 // CHECK-SAME: (%[[SRC4:.*]]: !pto.ptr, %[[OFF4:.*]]: index) -// CHECK: %[[E2B4:.*]] = pto.vlds %[[SRC4]][{{.*}}] {dist = "E2B_B32"} : !pto.ptr -> !pto.vreg<64xf32> -// CHECK: return %[[E2B4]], %[[E2B4]] : !pto.vreg<64xf32>, !pto.vreg<64xf32> +// CHECK-NOT: E2B_B32 +// CHECK: pto.vsldb +// CHECK: pto.vselr +// CHECK: return {{%.*}}, {{%.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32> +// CHECK-NOT: E2B_B32 diff --git a/test/lit/vmi_new/vmi_pre_assignment_combine_group_slot_broadcast_load.pto b/test/lit/vmi_new/vmi_pre_assignment_combine_group_slot_broadcast_load.pto index 729ccfeb75..2a3e23f1fc 100644 --- a/test/lit/vmi_new/vmi_pre_assignment_combine_group_slot_broadcast_load.pto +++ b/test/lit/vmi_new/vmi_pre_assignment_combine_group_slot_broadcast_load.pto @@ -74,13 +74,19 @@ module { // LOWER: pto.vselr // LOWER-LABEL: func.func @e2b_deint4_candidate -// LOWER: %[[E2B:.*]] = pto.vlds %{{.*}}[%{{.*}}] {dist = "E2B_B32"} : !pto.ptr -> !pto.vreg<64xf32> -// LOWER-NOT: pto.vselr -// LOWER: return %[[E2B]], %[[E2B]], %[[E2B]], %[[E2B]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32> +// LOWER-NOT: E2B_B32 +// LOWER: pto.vsldb +// LOWER: pto.vselr +// LOWER: return {{%.*}}, {{%.*}}, {{%.*}}, {{%.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32> +// LOWER-NOT: E2B_B32 // LOWER-NOT: pto.vmi. // LOWER-LABEL: func.func @e2b_deint4_consumer -// LOWER: %[[E2B2:.*]] = pto.vlds %{{.*}}[%{{.*}}] {dist = "E2B_B32"} : !pto.ptr -> !pto.vreg<64xf32> -// LOWER-NOT: pto.vselr -// LOWER: pto.vmul %{{.*}}, %[[E2B2]] +// LOWER-NOT: E2B_B32 +// LOWER: pto.vsldb +// LOWER: pto.vselr +// LOWER: pto.vcvt {{%.*}}, {{%.*}} {part = "EVEN"} +// LOWER: pto.vmul {{%.*}}, {{%.*}} +// LOWER: return {{%.*}}, {{%.*}}, {{%.*}}, {{%.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.vreg<64xf32> +// LOWER-NOT: E2B_B32 // LOWER-NOT: pto.vmi.