[Mlir-commits] [mlir] [MLIR][XeGPU] add dpas_mx op definition and layout propagation rule (PR #194117)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Apr 24 22:58:32 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
This PR extends the DpasMx operation to support MXFP (microscaling floating point) matrix multiply with separate scale factor layouts.
1. Op Definition
Added layout_a_scale and layout_b_scale attributes to DpasMx op
Removed AllElementTypesMatch<["a", "b"]> trait to allow different types for A/B with scales
2. Layout Infrastructure
setupDpasMxLayout(): Creates anchor layouts for all 5 operands (A, B, C/D, scale_a, scale_b)
Derives scale layouts from parent matrix layouts by dividing innermost dimension
Supports all layout kinds: Subgroup, InstData, Lane
Fix a bug in getupDpasSubgroupLayouts(): sg_data of A/B matrix should keep the full K dimension.
3. Layout Propagation
visitDpasMxOp(): Propagates layout attributes from op to operands during dataflow analysis
---
Patch is 76.17 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/194117.diff
10 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td (+3-2)
- (modified) mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h (+12)
- (modified) mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h (+194-3)
- (modified) mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h (+11-6)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp (+315-112)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+120)
- (modified) mlir/test/Dialect/XeGPU/ops.mlir (+3-3)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+91-1)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir (+97-15)
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+85-1)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 31fe93d209a6d..8b6763e234091 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1554,7 +1554,6 @@ def XeGPU_TruncfOp
}
def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
- AllElementTypesMatch<["a", "b"]>,
AnchorLayoutInterface]> {
let summary = "It performs scaled mma computation";
@@ -1601,7 +1600,9 @@ def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
VectorOfRankAndType<[1, 2], [F8E8M0FNU]>]>>:$scale_b,
OptionalAttr<DistributeLayoutAttr>:$layout_a,
OptionalAttr<DistributeLayoutAttr>:$layout_b,
- OptionalAttr<DistributeLayoutAttr>:$layout_cd);
+ OptionalAttr<DistributeLayoutAttr>:$layout_cd,
+ OptionalAttr<DistributeLayoutAttr>:$layout_a_scale,
+ OptionalAttr<DistributeLayoutAttr>:$layout_b_scale);
let results = (outs XeGPU_DpasResType:$result);
let extraClassDeclaration = [{
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 83eb939cf1bec..5defec222347f 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -199,6 +199,18 @@ setupDpasLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
VectorType cdTy, DistributeLayoutAttr consumerLayout, int numSg,
const uArch::uArch *uArch);
+/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg layout
+/// creation. A_scale and B_scale are optional.
+std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr>>
+setupDpasMxLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
+ VectorType cdTy, std::optional<VectorType> aScaleTy,
+ std::optional<VectorType> bScaleTy,
+ DistributeLayoutAttr consumerLayout, int numSg,
+ const uArch::uArch *uArch);
+
/// Gets the expected layout for a given consumer operand. This will check if
/// the owning operation of the consumer operand is one of the special layout
/// users and determine the expected layout accordingly.
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
index 595965d414840..eca7ffdb890ec 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h
@@ -209,6 +209,56 @@ struct SubgroupMatrixMultiplyAcc : public Instruction,
unsigned getPackedFormatBitSizeA() const { return packedFormatBitSizeA; }
unsigned getPackedFormatBitSizeB() const { return packedFormatBitSizeB; }
+ bool isLaneLayoutRowMajorOrder() const override { return true; }
+
+protected:
+ const unsigned packedFormatBitSizeA;
+ const unsigned packedFormatBitSizeB;
+};
+
+struct SubgroupScaledMatrixMultiplyAcc : public Instruction,
+ public MMAInstructionInterface {
+ SubgroupScaledMatrixMultiplyAcc(unsigned packedFormatBitSizeA,
+ unsigned packedFormatBitSizeB)
+ : Instruction(InstructionKind::SubgroupScaledMatrixMultiplyAcc,
+ InstructionScope::Subgroup),
+ packedFormatBitSizeA(packedFormatBitSizeA),
+ packedFormatBitSizeB(packedFormatBitSizeB) {}
+ static bool classof(const Instruction *B) {
+ return B->getInstructionKind() ==
+ InstructionKind::SubgroupScaledMatrixMultiplyAcc;
+ }
+ // Source:
+ // https://github.com/intel/llvm/blob/sycl/sycl/doc/design/spirv-extensions/SPV_INTEL_subgroup_scaled_matrix_multiply_accumulate.asciidoc
+
+ // Override all virtuals from MatrixOpInterface
+ virtual llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+ getSupportedShapes(Type dataType, MMAOpndKind matrixType) override;
+ virtual llvm::SmallVector<Type, 8>
+ getSupportedTypes(MLIRContext &context, MMAOpndKind matrixType) override;
+ virtual bool
+ checkSupportedShapesAndTypes(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual bool checkSupportedTypes(Type AType, Type BType, Type CType,
+ Type DType) override;
+ virtual bool validate(std::pair<uint32_t, uint32_t> AShape,
+ std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape,
+ std::pair<uint32_t, uint32_t> DShape, Type AType,
+ Type BType, Type CType, Type DType) override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedM(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedK(Type type) const override;
+ virtual llvm::SmallVector<uint32_t, 8>
+ getSupportedN(Type type) const override;
+
+ unsigned getPackedFormatBitSizeA() const { return packedFormatBitSizeA; }
+ unsigned getPackedFormatBitSizeB() const { return packedFormatBitSizeB; }
+ bool isLaneLayoutRowMajorOrder() const override { return true; }
protected:
const unsigned packedFormatBitSizeA;
@@ -282,14 +332,15 @@ struct BMGuArch : public Xe2Plus {
struct CRIuArch : public Xe2Plus {
static llvm::ArrayRef<const Instruction *> getInstructionRegistryArr() {
static const SubgroupMatrixMultiplyAcc dpasInst{16, 32};
+ static const SubgroupScaledMatrixMultiplyAcc dpasMxInst{16, 32};
static const Subgroup2DBlockLoadInstruction loadNdInst;
static const Subgroup2DBlockStoreInstruction storeNdInst;
static const Subgroup2DBlockPrefetchInstruction prefetchNdInst;
static const SpirvStoreScatterInstruction storeScatterInst;
static const SpirvLoadGatherInstruction loadGatherInst;
- static const Instruction *arr[] = {&dpasInst, &loadNdInst,
- &storeNdInst, &prefetchNdInst,
- &storeScatterInst, &loadGatherInst};
+ static const Instruction *arr[] = {
+ &dpasInst, &dpasMxInst, &loadNdInst, &storeNdInst,
+ &prefetchNdInst, &storeScatterInst, &loadGatherInst};
return arr;
}
@@ -477,4 +528,144 @@ SubgroupMatrixMultiplyAcc::getSupportedN(Type type) const {
return {16};
}
+//===----------------------------------------------------------------------===//
+// SubgroupScaledMatrixMultiplyAcc implementations
+//===----------------------------------------------------------------------===//
+
+inline llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16>
+SubgroupScaledMatrixMultiplyAcc::getSupportedShapes(Type dataType,
+ MMAOpndKind matrixType) {
+ auto combineVectors = [](const llvm::SmallVector<uint32_t, 8> &a,
+ const llvm::SmallVector<uint32_t, 8> &b)
+ -> llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> {
+ llvm::SmallVector<std::pair<uint32_t, uint32_t>, 16> result;
+ for (unsigned x : a) {
+ for (unsigned y : b) {
+ result.emplace_back(x, y);
+ }
+ }
+ return result;
+ };
+
+ // Avoid calling getSupportedK for C/D types (which are f32/bf16
+ // and not valid for the K-dimension bit-width calculation).
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return combineVectors(getSupportedM(dataType), getSupportedK(dataType));
+ case MMAOpndKind::MatrixB:
+ return combineVectors(getSupportedK(dataType), getSupportedN(dataType));
+ case MMAOpndKind::MatrixC:
+ case MMAOpndKind::MatrixD:
+ return combineVectors(getSupportedM(dataType), getSupportedN(dataType));
+ }
+ return {};
+}
+
+inline llvm::SmallVector<Type, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedTypes(MLIRContext &context,
+ MMAOpndKind matrixType) {
+ Type f8E4M3FNType = Float8E4M3FNType::get(&context);
+ Type f8E5M2Type = Float8E5M2Type::get(&context);
+ Type f4E2M1FNType = Float4E2M1FNType::get(&context);
+ Type bf16Type = BFloat16Type::get(&context);
+ Type f32Type = Float32Type::get(&context);
+
+ switch (matrixType) {
+ case MMAOpndKind::MatrixA:
+ return {f8E4M3FNType, f8E5M2Type, f4E2M1FNType};
+ case MMAOpndKind::MatrixB:
+ return {f8E4M3FNType, f8E5M2Type, f4E2M1FNType};
+ case MMAOpndKind::MatrixC:
+ return {bf16Type, f32Type};
+ case MMAOpndKind::MatrixD:
+ return {bf16Type, f32Type};
+ }
+ return {};
+}
+
+inline bool SubgroupScaledMatrixMultiplyAcc::checkSupportedTypes(Type AType,
+ Type BType,
+ Type CType,
+ Type DType) {
+ auto isSupportedLowPrecision = [](Type t) {
+ return t.isF8E4M3FN() || t.isF8E5M2() || llvm::isa<Float4E2M1FNType>(t);
+ };
+ auto isSupportedAccum = [](Type t) { return t.isF32() || t.isBF16(); };
+
+ if (!isSupportedLowPrecision(AType) || !isSupportedLowPrecision(BType)) {
+ LDBG() << "Unsupported scaled dpas: A and B must be FP8 or FP4 types.";
+ return false;
+ }
+
+ // A and B must have the same bit width for K dimension compatibility.
+ if (AType.getIntOrFloatBitWidth() != BType.getIntOrFloatBitWidth()) {
+ LDBG() << "Unsupported scaled dpas: A and B must have the same bit width.";
+ return false;
+ }
+
+ if (CType && !isSupportedAccum(CType)) {
+ LDBG() << "Unsupported scaled dpas: C must be f32 or bf16.";
+ return false;
+ }
+
+ if (!isSupportedAccum(DType)) {
+ LDBG() << "Unsupported scaled dpas: D must be f32 or bf16.";
+ return false;
+ }
+
+ return true;
+}
+
+inline bool SubgroupScaledMatrixMultiplyAcc::checkSupportedShapesAndTypes(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ auto supportedAShapes = getSupportedShapes(AType, MMAOpndKind::MatrixA);
+ auto supportedBShapes = getSupportedShapes(BType, MMAOpndKind::MatrixB);
+ auto supportedCShapes = getSupportedShapes(CType, MMAOpndKind::MatrixC);
+ auto supportedDShapes = getSupportedShapes(DType, MMAOpndKind::MatrixD);
+ return llvm::is_contained(supportedAShapes, AShape) &&
+ llvm::is_contained(supportedBShapes, BShape) &&
+ llvm::is_contained(supportedCShapes, CShape) &&
+ llvm::is_contained(supportedDShapes, DShape) &&
+ checkSupportedTypes(AType, BType, CType, DType);
+}
+
+inline bool SubgroupScaledMatrixMultiplyAcc::validate(
+ std::pair<uint32_t, uint32_t> AShape, std::pair<uint32_t, uint32_t> BShape,
+ std::pair<uint32_t, uint32_t> CShape, std::pair<uint32_t, uint32_t> DShape,
+ Type AType, Type BType, Type CType, Type DType) {
+ return checkSupportedShapesAndTypes(AShape, BShape, CShape, DShape, AType,
+ BType, CType, DType);
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedM(Type type) const {
+ return {8};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedK(Type type) const {
+ assert(type.isIntOrFloat() && "Matrix type must be int or float");
+ auto bitWidth = type.getIntOrFloatBitWidth();
+ uint32_t kSize = 0;
+ switch (bitWidth) {
+ case 4:
+ kSize = 64; // FP4: scale K by 4 (base 16-bit K=16 -> 64)
+ break;
+ case 8:
+ kSize = 32; // FP8: scale K by 2 (base 16-bit K=16 -> 32)
+ break;
+ default:
+ llvm_unreachable("Scaled dpas only supports FP8 (8-bit) and FP4 (4-bit) "
+ "types for A/B matrices");
+ }
+ return {kSize};
+}
+
+inline llvm::SmallVector<uint32_t, 8>
+SubgroupScaledMatrixMultiplyAcc::getSupportedN(Type type) const {
+ return {16};
+}
+
#endif // MLIR_DIALECT_XEGPU_UARCH_INTELGPUXE2_H
diff --git a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
index 0f9d052e11147..147a56a52c188 100644
--- a/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
+++ b/mlir/include/mlir/Dialect/XeGPU/uArch/uArchBase.h
@@ -36,11 +36,14 @@ enum class InstructionScope { Lane, Subgroup, Workgroup, Cluster };
enum class InstructionKind {
SubgroupMatrixMultiplyAcc, // Dot Product Accumulate Systolic (DPAS) is a
// matrix multiply-add operation
- Subgroup2DBlockStore, // Subgroup-level 2D block write instruction
- Subgroup2DBlockLoad, // Subgroup-level 2D block load instruction
- Subgroup2DBlockPrefetch, // Subgroup-level 2D block prefetch instruction
- StoreScatter, // Lane-level store (scalar, vector)
- LoadGather, // Lane-level load (scalar, vector)
+ SubgroupScaledMatrixMultiplyAcc, // Scaled Matrix Multiply Accumulate is a
+ // DPAS with scaling factor applied to
+ // operand A or B before multiplication
+ Subgroup2DBlockStore, // Subgroup-level 2D block write instruction
+ Subgroup2DBlockLoad, // Subgroup-level 2D block load instruction
+ Subgroup2DBlockPrefetch, // Subgroup-level 2D block prefetch instruction
+ StoreScatter, // Lane-level store (scalar, vector)
+ LoadGather, // Lane-level load (scalar, vector)
// @TODO: Add more instructions as needed
};
@@ -61,6 +64,8 @@ struct Instruction {
switch (instKind) {
case InstructionKind::SubgroupMatrixMultiplyAcc:
return "dpas";
+ case InstructionKind::SubgroupScaledMatrixMultiplyAcc:
+ return "dpas_mx";
case InstructionKind::Subgroup2DBlockStore:
return "store_nd";
case InstructionKind::Subgroup2DBlockLoad:
@@ -246,7 +251,7 @@ struct MMAInstructionInterface {
virtual llvm::SmallVector<uint32_t, 8> getSupportedM(Type type) const = 0;
virtual llvm::SmallVector<uint32_t, 8> getSupportedK(Type type) const = 0;
virtual llvm::SmallVector<uint32_t, 8> getSupportedN(Type type) const = 0;
-
+ virtual bool isLaneLayoutRowMajorOrder() const = 0;
virtual ~MMAInstructionInterface() = default;
};
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7d48315eec6ff..518f7b181ff5d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1179,6 +1179,136 @@ getValidLayouts(ArrayRef<int64_t> wgShape, ArrayRef<int64_t> instData,
return candidates;
}
+/// Helper function to compute inst_data vectors for DPAS operands A, B, and
+/// C/D.
+static std::optional<std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>>
+getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
+ const xegpu::uArch::uArch *uArch,
+ bool isDpasMx = false) {
+ const int subgroupSize = uArch->getSubgroupSize();
+
+ const xegpu::uArch::MMAInstructionInterface *uArchInstruction;
+ if (isDpasMx)
+ uArchInstruction = dyn_cast<xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
+ else
+ uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+
+ const unsigned dataALen = aTy.getShape().front();
+ auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());
+ const int maxALen =
+ xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));
+
+ const unsigned dataBLen = bTy.getShape().back();
+ auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());
+ const int maxBLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));
+
+ auto supportedCLen = uArchInstruction->getSupportedN(cdTy.getElementType());
+ const int maxCLen =
+ xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedCLen));
+ if (maxALen == -1 || maxBLen == -1 || maxCLen == -1)
+ return std::nullopt;
+
+ // For DPAS_MX, use getSupportedK to get the scaled K dimension
+ int kDimSize = subgroupSize;
+ if (isDpasMx) {
+ auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (!supportedKLen.empty())
+ kDimSize = supportedKLen[0];
+ }
+
+ SmallVector<int64_t> instDataA(aTy.getRank(), 1);
+ instDataA[aTy.getRank() - 2] = maxALen;
+ instDataA[aTy.getRank() - 1] = kDimSize;
+ SmallVector<int64_t> instDataB(bTy.getRank(), 1);
+ instDataB[bTy.getRank() - 2] = kDimSize;
+ instDataB[bTy.getRank() - 1] = maxBLen;
+ SmallVector<int64_t> instDataCD(cdTy.getRank(), 1);
+ instDataCD[cdTy.getRank() - 2] = maxALen;
+ instDataCD[cdTy.getRank() - 1] = maxCLen;
+ return std::make_tuple(instDataA, instDataB, instDataCD);
+}
+
+/// Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
+/// Returns the three layouts if successful, nullopt otherwise.
+static std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+getupDpasSubgroupLayouts(mlir::MLIRContext *context, VectorType aTy,
+ VectorType bTy, VectorType cdTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
+ if (!instDataVecs)
+ return std::nullopt;
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+ assert(instDataA.size() == 2 && instDataB.size() == 2 &&
+ instDataCD.size() == 2 &&
+ "Sg layout creation expects valid 2D inst data");
+
+ std::optional<LayoutRepresentation> consumerSgLayout = std::nullopt;
+ if (consumerLayout && consumerLayout.isForWorkgroup()) {
+ SmallVector<int64_t> sgLayoutD = consumerLayout.getEffectiveSgLayoutAsInt();
+ consumerSgLayout = std::make_pair(sgLayoutD[0], sgLayoutD[1]);
+ }
+
+ // Get all valid layouts for A, B and C/D operands
+ auto layoutsA = getValidLayouts(aTy.getShape(), instDataA, numSg);
+ auto layoutsB = getValidLayouts(bTy.getShape(), instDataB, numSg);
+ auto layoutsCD = getValidLayouts(cdTy.getShape(), instDataCD, numSg);
+ if (layoutsA.empty() || layoutsB.empty() || layoutsCD.empty())
+ return std::nullopt;
+
+ // Pick the best subgroup layout
+ llvm::DenseSet<LayoutRepresentation> setA(layoutsA.begin(), layoutsA.end());
+ llvm::DenseSet<LayoutRepresentation> setCD(layoutsCD.begin(),
+ layoutsCD.end());
+ std::optional<LayoutRepresentation> bestPick;
+ for (auto &sgLayout : layoutsB) {
+ if (setA.contains(sgLayout) && setCD.contains(sgLayout)) {
+ if (consumerSgLayout.has_value() && sgLayout == *consumerSgLayout) {
+ bestPick = sgLayout;
+ break;
+ }
+ if (!bestPick)
+ bestPick = sgLayout;
+ }
+ }
+ if (!bestPick)
+ return std::nullopt;
+
+ SmallVector<int> sgLayout = {static_cast<int>(bestPick->first),
+ static_cast<int>(bestPick->second)};
+ SmallVector<int> sgDataA = {static_cast<int>(aTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(aTy.getShape()[1])};
+ SmallVector<int> sgDataB = {
+ static_cast<int>(bTy.getShape()[0]),
+ static_cast<int>(bTy.getShape()[1] / sgLayout[1])};
+ SmallVector<int> sgDataCD = {
+ static_cast<int>(cdTy.getShape()[0] / sgLayout[0]),
+ static_cast<int>(cdTy.getShape()[1] / sgLayout[1])};
+
+ auto dpasALayout =
...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/194117
More information about the Mlir-commits
mailing list