diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index 186e0b7eff..36ccdda71b 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -412,6 +412,7 @@ struct LayoutSolver { if (failed(fact)) { return {}; } + return fact->layout.resultLayout; } diff --git a/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc b/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc index bc23a7964b..4663b86667 100644 --- a/lib/PTO/Transforms/VMILayoutSupportPatternDSL.inc +++ b/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(sourceType) && isa(resultType); + case CastTypeClass::F4E2M1x2ToBF16x2: + return isa(sourceType) && + 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..a654845054 100644 --- a/lib/PTO/Transforms/VMILayoutSupportTables.inc +++ b/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)}, diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 75873dcb7a..84ffb9597b 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,43 @@ 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; + 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(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");