[WIP]solve fp4 -> bf16 layout

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<int64_t> contiguousArity = getVMIPhysicalArity(contiguousType);
+      if (succeeded(contiguousArity) && *contiguousArity > 1) {
+        FailureOr<
+            SmallVector<VMIGroupBroadcastLayoutFact, mlir::pto::kValue4>>
+            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<VMIDeinterleaveLoadOp>([this, op](auto load) {
           return addDeinterleaveLoadConstraint(load, op);
         })
+        .Case<VMILoadOp>([this, op](auto load) {
+          auto type = cast<VMIVRegType>(load.getResult().getType());
+          FailureOr<int64_t> 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<WalkResult>(std::nullopt);
+          }
+          int64_t laneStride = *lanesPerPart / type.getElementCount();
+          return constraintResult(setNaturalLayout(
+              load.getResult(), VMILayoutAttr::getContiguous(ctx, laneStride),
+              op));
+        })
         .Case<VMIMaskedLoadOp, VMIExpandLoadOp>([this, op](auto load) {
           requestDataUse(load.getPassthruMutable(), getContiguousLayout());
           return constraintResult(setNaturalLayout(

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<IntegerType>(sourceType) && isa<IntegerType>(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");
 }

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)},

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<int64_t> 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<VMIExtFOp> {
 
     ArrayRef<StringRef> 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<StringRef>(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<StringRef>(kPacked4Parts, factor);
     } else {
       return rewriter.notifyMatchFailure(
           op, "unsupported physical extf source/result width relation");
@@ -11327,15 +11370,15 @@ struct OneToNVMIExtFOpPattern : OneToNOpConversionPattern<VMIExtFOp> {
 
     SmallVector<Value> 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();

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<ve
 }
 
 // CHECK-LABEL: func.func @ComputeY1ToFP8_fp16_e4m3_vmi
-// CHECK: pto.vlds {{.*}} {dist = "E2B_B16"} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
+// CHECK-NOT: E2B_B16
+// CHECK: pto.vsldb
+// CHECK: pto.vselr
 // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
 // CHECK: scf.for
-// CHECK: pto.vldsx2 {{.*}} "DINTLV_B16" : !pto.ptr<f16, ub>, index -> !pto.vreg<128xf16>, !pto.vreg<128xf16>
+// CHECK: pto.vlds {{.*}} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
+// CHECK: pto.vlds {{.*}} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
 // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
 // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
 // CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
 // CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !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<f8E4M3FN, ub>, !pto.mask<b8>
+// CHECK: pto.vsts {{.*}} {dist = "PK_B16"} : !pto.vreg<256xf8E4M3FN>, !pto.ptr<f8E4M3FN, ub>, !pto.mask<b16>
+// CHECK-NOT: E2B_B16
 
 // CHECK-LABEL: func.func @ComputeY1ToFP8_fp16_e5m2_vmi
-// CHECK: pto.vlds {{.*}} {dist = "E2B_B16"} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
+// CHECK-NOT: E2B_B16
+// CHECK: pto.vsldb
+// CHECK: pto.vselr
 // CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
 // CHECK: scf.for
-// CHECK: pto.vldsx2 {{.*}} "DINTLV_B16" : !pto.ptr<f16, ub>, index -> !pto.vreg<128xf16>, !pto.vreg<128xf16>
+// CHECK: pto.vlds {{.*}} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
+// CHECK: pto.vlds {{.*}} : !pto.ptr<f16, ub> -> !pto.vreg<128xf16>
+// CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
+// CHECK: pto.vcvt {{.*}} {part = "EVEN"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
+// CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !pto.vreg<64xf32>
+// CHECK: pto.vcvt {{.*}} {part = "ODD"} : !pto.vreg<128xf16>, !pto.mask<b16> -> !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<f8E5M2, ub>, !pto.mask<b8>
+// CHECK: pto.vsts {{.*}} {dist = "PK_B16"} : !pto.vreg<256xf8E5M2>, !pto.ptr<f8E5M2, ub>, !pto.mask<b16>
+// CHECK-NOT: E2B_B16

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<bf16, ub>, %[[OFF2:.*]]: index)
-// CHECK: %[[E2B2:.*]] = pto.vlds %[[SRC2]][{{.*}}] {dist = "E2B_B16"} : !pto.ptr<bf16, ub> -> !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<f32, ub>, %[[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<f32, ub>, %[[OFF4:.*]]: index)
-// CHECK: %[[E2B4:.*]] = pto.vlds %[[SRC4]][{{.*}}] {dist = "E2B_B32"} : !pto.ptr<f32, ub> -> !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

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<f32, ub> -> !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<f32, ub> -> !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.