improve VMI FP4 dequant support

lib/PTO/Transforms/VMILayoutAssignment.cpp

@@ -412,6 +412,7 @@ struct LayoutSolver {
     if (failed(fact)) {
       return {};
     }
+
     return fact->layout.resultLayout;
   }
 

lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc

@@ -28,6 +28,8 @@ enum class CastTypeClass {
   Float,
   // Both source and result are integer element types.
   Integer,
+  // Packed E2M1 pairs widened to the register-only BF16-pair carrier.
+  F4E2M1x2ToBF16x2,
 };
 
 struct LayoutPattern {
@@ -138,6 +140,9 @@ static bool matchesCastTypeClass(CastTypeClass typeClass, Type sourceType,
             pto::isPTOLowPrecisionType(resultType));
   case CastTypeClass::Integer:
     return isa<IntegerType>(sourceType) && isa<IntegerType>(resultType);
+  case CastTypeClass::F4E2M1x2ToBF16x2:
+    return isa<pto::F4E2M1x2Type>(sourceType) &&
+           pto::isPTOBF16x2Type(resultType);
   }
   llvm_unreachable("unknown cast type class");
 }

lib/PTO/Transforms/VMILayoutSupportTables.inc

@@ -183,6 +183,8 @@ 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()},
+    {bits<8>(), bits<32>(), 128, ls(4), c(),
+     CastTypeClass::F4E2M1x2ToBF16x2},
     {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)},

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,43 @@ 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;
+      if (resultTypes.size() != 2 * sourceParts.size()) {
+        return rewriter.notifyMatchFailure(
+            op, "unsupported physical extf source/result width relation");
+      }
       factor = 2;
-    } else if (sourceBits == 8 &&
-               resultTypes.size() == 4 * sourceParts.size()) {
+      parts = kEvenOddParts;
+    } else if (sourceBits == 8) {
       static constexpr StringRef kPacked4Parts[] = {"P0", "P1", "P2", "P3"};
-      parts = kPacked4Parts;
-      factor = 4;
+      static constexpr StringRef kPacked2Parts[] = {"P0", "P2"};
+      static constexpr StringRef kPacked1Parts[] = {"P0"};
+      bool isPackedE2M1 =
+          isa<pto::F4E2M1x2Type>(sourceVMIType.getElementType());
+      if (!isPackedE2M1 && resultTypes.size() != 4 * sourceParts.size()) {
+        return rewriter.notifyMatchFailure(
+            op, "unsupported physical extf source/result width relation");
+      }
+      if (!isPackedE2M1) {
+        factor = 4;
+        parts = kPacked4Parts;
+      } else if (sourceLayout && sourceLayout.isContiguous() &&
+                 sourceLayout.getLaneStride() == 4) {
+        factor = 1;
+        parts = kPacked1Parts;
+      } else if (sourceLayout && sourceLayout.isContiguous() &&
+                 sourceLayout.getLaneStride() == 2) {
+        factor = 2;
+        parts = kPacked2Parts;
+      } else {
+        factor = 4;
+        parts = kPacked4Parts;
+      }
+      if (resultTypes.size() != factor * sourceParts.size()) {
+        return rewriter.notifyMatchFailure(
+            op, "FP4 extf physical arity does not match its source layout");
+      }
     } else {
       return rewriter.notifyMatchFailure(
           op, "unsupported physical extf source/result width relation");