[Mlir-commits] [mlir] [mlir][xegpu] add support for structure control flow ops in workgroup to subgroup distribution (PR #142618)
Chao Chen
llvmlistbot at llvm.org
Tue Jun 3 07:48:49 PDT 2025
https://github.com/chencha3 created https://github.com/llvm/llvm-project/pull/142618
None
>From a808983eac590930a17863152cb2a6f2aa855141 Mon Sep 17 00:00:00 2001
From: Chao Chen <chao.chen at intel.com>
Date: Tue, 3 Jun 2025 14:44:03 +0000
Subject: [PATCH] add scf support
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 3 +
.../Transforms/XeGPUWgToSgDistribute.cpp | 106 ++++++++++++++++--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 4 +-
.../test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir | 2 +-
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 6 +-
5 files changed, 104 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index f9327d63869c0..6fea10185402a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -26,6 +26,9 @@ class TensorDescType;
namespace xegpu {
+/// Flatten a set of ValueRange into a single SmallVector<Value>
+SmallVector<Value> flattenValues(ArrayRef<ValueRange> values);
+
/// If tensor descriptor has a layout attribute it is used in SIMT mode.
/// In this mode, the distributed vector shape is determined as follows:
/// Definitions:
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 3bf76af674ba0..ad12cf34ca7b3 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
+#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Transforms/DialectConversion.h"
namespace mlir {
@@ -29,6 +30,30 @@ using namespace mlir;
namespace {
+static std::pair<SmallVector<int64_t>, int>
+computeTileShapeAndCount(ArrayRef<int64_t> shape, xegpu::LayoutAttr layout) {
+ // init count and subShape to the default value. If the LayoutAttr
+ // is not present, it will return a VectorType with original shape.
+ int count = 1;
+ SmallVector<int64_t> tileShape(shape);
+
+ if (layout) {
+ if (DenseI32ArrayAttr sgLayoutAttr = layout.getSgLayout()) {
+ auto sgLayout = llvm::to_vector_of<int64_t>(sgLayoutAttr.asArrayRef());
+ if (DenseI32ArrayAttr sgDataAttr = layout.getSgData())
+ tileShape = llvm::to_vector_of<int64_t>(sgDataAttr.asArrayRef());
+ else
+ tileShape = computeShapeRatio(shape, sgLayout).value_or(tileShape);
+ SmallVector<int64_t> distUnit =
+ computeElementwiseMul(sgLayout, tileShape);
+ for (size_t i = 0; i < distUnit.size(); ++i)
+ distUnit[i] = std::min(shape[i], distUnit[i]);
+ count = computeProduct(shape) / computeProduct(distUnit);
+ }
+ }
+ return std::make_pair(tileShape, count);
+}
+
/// This pattern transforms the CreateNdDescOp to create a subgroup descriptor
/// from a workgroup descriptor. It replaces the offsets and sizes with
/// appropriate values for the subgroup.
@@ -266,15 +291,15 @@ struct WgToSgDpasOp : public OpConversionPattern<xegpu::DpasOp> {
if (resultTy.getRank() != 2)
return failure();
- auto originalLayout =
- llvm::dyn_cast_or_null<xegpu::LayoutAttr>(op->getAttr("layout"));
+ auto originalLayout = xegpu::getLayoutAttr(op.getResult());
if (!originalLayout)
return failure();
- SmallVector<Value> newDpasOps;
size_t i = 0;
+ SmallVector<Value> newDpasOps;
for (auto aVec : adaptor.getLhs()) {
for (auto bVec : adaptor.getRhs()) {
+
llvm::SmallVector<Value> operands({aVec, bVec});
Value tmpC;
if (op.getAcc()) {
@@ -288,10 +313,9 @@ struct WgToSgDpasOp : public OpConversionPattern<xegpu::DpasOp> {
llvm::cast<VectorType>(bVec.getType()).getShape();
VectorType resTy = VectorType::get({aVecShape[0], bVecShape[1]},
resultTy.getElementType());
- tmpC = rewriter.create<xegpu::DpasOp>(
- loc, resTy, operands,
- llvm::ArrayRef<NamedAttribute>(
- {"layout_result_0", originalLayout.dropSgLayoutAndData()}));
+ tmpC = rewriter.create<xegpu::DpasOp>(loc, resTy, operands);
+ xegpu::setLayoutAttr(cast<OpResult>(tmpC), originalLayout.dropSgLayoutAndData());
+
newDpasOps.push_back(tmpC);
}
}
@@ -314,14 +338,30 @@ struct WgToSgPrefetchNdOp : public OpConversionPattern<xegpu::PrefetchNdOp> {
}
};
+struct UnrealizedConversionCastOpPattern
+ : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
+ using OpConversionPattern<
+ mlir::UnrealizedConversionCastOp>::OpConversionPattern;
+
+ mlir::LogicalResult
+ matchAndRewrite(mlir::UnrealizedConversionCastOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (op.getNumOperands() == 1 && op.getNumResults() == 1) {
+ rewriter.replaceOpWithMultiple(op, xegpu::flattenValues(adaptor.getInputs()));
+ return mlir::success();
+ }
+ return mlir::failure();
+ }
+};
+
} // namespace
namespace mlir {
namespace xegpu {
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
patterns.add<WgToSgCreateNdOp, WgToSgLoadNdOp, WgToSgStoreNdOp,
- WgToSgUpdateNdOffsetOp, WgToSgDpasOp, WgToSgPrefetchNdOp>(
- patterns.getContext());
+ WgToSgUpdateNdOffsetOp, WgToSgDpasOp, WgToSgPrefetchNdOp,
+ UnrealizedConversionCastOpPattern>(patterns.getContext());
}
} // namespace xegpu
} // namespace mlir
@@ -353,7 +393,7 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
};
auto isLegal = [&](xegpu::LayoutAttr layout) -> bool {
- return !layout || layout.getSgLayout() == nullptr;
+ return !layout || !layout.isWgLayout();
};
target.addDynamicallyLegalOp<xegpu::CreateNdDescOp, xegpu::LoadNdOp,
@@ -365,12 +405,56 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
});
target.addDynamicallyLegalOp<xegpu::DpasOp>([=](xegpu::DpasOp op) -> bool {
- auto layout = dyn_cast_or_null<xegpu::LayoutAttr>(op->getAttr("layout"));
+ auto layout = xegpu::getLayoutAttr(op.getResult());
return isLegal(layout);
});
+ target.addIllegalOp<UnrealizedConversionCastOp>();
+
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+ TypeConverter converter;
+ converter.addConversion([&](Type type) -> Type { return type; });
+ converter.addConversion(
+ [&](RankedTensorType type,
+ SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+ Type elemTy = type.getElementType();
+ ArrayRef<int64_t> shape = type.getShape();
+
+ int count;
+ SmallVector<int64_t> subShape;
+ std::tie(subShape, count) = computeTileShapeAndCount(
+ shape, dyn_cast<xegpu::LayoutAttr>(type.getEncoding()));
+
+ auto newTy = VectorType::get(subShape, elemTy);
+ result.append(count, newTy);
+ return success();
+ });
+
+ converter.addConversion(
+ [&](xegpu::TensorDescType type,
+ SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+ Type elemTy = type.getElementType();
+ ArrayRef<int64_t> shape = type.getShape();
+
+ // init count and newTy to the default value. If the layout
+ // attribute is not present, it will return the original type.
+ int count;
+ SmallVector<int64_t> subShape;
+ xegpu::LayoutAttr layout = type.getLayoutAttr();
+ std::tie(subShape, count) = computeTileShapeAndCount(shape, layout);
+
+ if (layout)
+ layout = layout.dropSgLayoutAndData();
+
+ auto newTy = xegpu::TensorDescType::get(
+ type.getContext(), subShape, elemTy, type.getEncoding(), layout);
+ result.append(count, newTy);
+ return success();
+ });
+
+ xegpu::doSCFStructuralTypeConversionWithTensorType(getOperation(), converter);
+
xegpu::populateXeGPUWgToSgDistributePatterns(patterns);
if (failed(
applyPartialConversion(getOperation(), target, std::move(patterns))))
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index dcaf4e85a82c5..ed48e3cc13117 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -27,7 +27,7 @@
using namespace mlir;
/// convert ArrayRef<ValueRange> into SmallVector<Value>
-static SmallVector<Value> flattenValues(ArrayRef<ValueRange> values) {
+SmallVector<Value> xegpu::flattenValues(ArrayRef<ValueRange> values) {
SmallVector<Value> result;
for (const auto &vals : values)
llvm::append_range(result, vals);
@@ -342,7 +342,7 @@ void xegpu::doSCFStructuralTypeConversionWithTensorType(
}
if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {
- SmallVector<Value> values = flattenValues(adaptor.getInputs());
+ SmallVector<Value> values = xegpu::flattenValues(adaptor.getInputs());
auto newOp = rewriter.create<UnrealizedConversionCastOp>(
op.getLoc(), outputTy, values);
rewriter.replaceOp(op, newOp);
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
index bee026eb2084d..fa1e5fbae0954 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
@@ -85,7 +85,7 @@ gpu.module @test_round_robin_assignment {
%tdesc_c = xegpu.create_nd_tdesc %c[0, 0] : memref<8x8xf32>
-> !xegpu.tensor_desc<8x8xf32, #xegpu.layout<sg_layout = [2, 2], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>
%dpas = xegpu.dpas %load_a, %load_b
- {layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
+ {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
: vector<8x8xf32>, vector<8x8xf32> -> vector<8x8xf32>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 7e89ada934071..22374f74b133e 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -5,7 +5,7 @@
gpu.module @test_1_1_assignment {
// CHECK-LABEL: test_create_nd_tdesc
// CHECK-SAME: %[[ARG_0:.*]]: memref<24x32xf32>
- gpu.func @test_create_nd_tdesc(%src: memref<24x32xf32>) {
+ gpu.func @test_create_nd_tdesc(%src: memref<24x32xf32>) {
// CHECK: %[[SGID:.*]] = gpu.subgroup_id
// CHECK: %[[C12:.*]] = arith.constant 12 : index
// CHECK: %[[C4:.*]] = arith.constant 4 : index
@@ -108,7 +108,7 @@ gpu.func @test_dpas(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
: !xegpu.tensor_desc<32x24xf32, #xegpu.layout<sg_layout = [4, 2], sg_data = [8, 12], lane_layout = [8, 2], lane_data = [1, 1]>>
-> vector<32x24xf32>
%dpas = xegpu.dpas %load_a, %load_b
- {layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [12, 12], lane_layout = [2, 2], lane_data = [1, 1]>}
+ {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [12, 12], lane_layout = [2, 2], lane_data = [1, 1]>}
: vector<24x32xf32>, vector<32x24xf32> -> vector<24x24xf32>
gpu.return
}
@@ -142,7 +142,7 @@ gpu.func @test_dpas_no_sg_data(%a: memref<24x32xf32>, %b: memref<32x24xf32>) {
: !xegpu.tensor_desc<32x24xf32, #xegpu.layout<sg_layout = [4, 2], lane_layout = [8, 2], lane_data = [1, 1]>>
-> vector<32x24xf32>
%dpas = xegpu.dpas %load_a, %load_b
- {layout = #xegpu.layout<sg_layout = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
+ {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
: vector<24x32xf32>, vector<32x24xf32> -> vector<24x24xf32>
gpu.return
}
More information about the Mlir-commits
mailing list