[Mlir-commits] [mlir] a5a5929 - [MLIR][XeGPU] Optimize instruction granularity for elemwise (#220986)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 5 09:55:11 PDT 2026
Author: Artem Kroviakov
Date: 2026-10-05T16:55:03Z
New Revision: a5a5929cc0e95299c2923e8c9434876b869bb471
URL: https://github.com/llvm/llvm-project/commit/a5a5929cc0e95299c2923e8c9434876b869bb471
DIFF: https://github.com/llvm/llvm-project/commit/a5a5929cc0e95299c2923e8c9434876b869bb471.diff
LOG: [MLIR][XeGPU] Optimize instruction granularity for elemwise (#220986)
Added:
mlir/test/Dialect/XeGPU/sink-elementwise-conversions.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index a79de4983bc2e..8921380093e10 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -41,6 +41,10 @@ LogicalResult propagateLayouts(OpBuilder &builder, Operation *target,
LogicalResult resolveLayoutConflicts(Operation *target);
+/// Sink `xegpu.convert_layout` ops past the elementwise operations they feed,
+/// so that an elementwise op runs in the coarser layout.
+void sinkElementwiseConversions(OpBuilder &builder, Operation *target);
+
/// Callable returning the propagated layout for a given Value, used by the
/// layout-propagation helpers below.
using GetLayoutFnTy = llvm::function_ref<DistributeLayoutAttr(Value)>;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 0b5a773ffc4ff..3629e122b1bf9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -13,6 +13,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
@@ -33,6 +34,7 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
@@ -1639,12 +1641,15 @@ LogicalResult ResolveLayoutConflicts::run() {
return WalkResult::advance();
});
+ if (r.wasInterrupted())
+ return failure();
+
LLVM_DEBUG({
DBGS() << "IR after resolving layout conflicts:\n";
parentOp->dump();
});
- return r.wasInterrupted() ? failure() : success();
+ return success();
}
LogicalResult ResolveLayoutConflicts::assignResultLayout(OpResult &result) {
@@ -1830,6 +1835,123 @@ static LogicalResult updateOpWithForwardFill(mlir::OpBuilder &builder,
return success();
}
+/// Optimize elementwise operations by sinking costly layout conversion.
+///
+/// %m = vector.create_mask ... {L1}
+/// %d = arith.mulf ... {L1}
+/// %cm = xegpu.convert_layout %m : L1 -> L2
+/// %cd = xegpu.convert_layout %d : L1 -> L2
+/// %s = arith.select %cm, %cd, %splat {L2}
+///
+/// becomes
+///
+/// %s' = arith.select %m, %d, %splat {L1}
+/// %s = xegpu.convert_layout %s' : L1 -> L2
+///
+/// We only sink if a coarser layout applies to an elementwise op
+///
+void xegpu::sinkElementwiseConversions(OpBuilder &builder,
+ Operation *parentOp) {
+ llvm::SmallSetVector<xegpu::ConvertLayoutOp, 8> deadConverts;
+ parentOp->walk([&](Operation *op) {
+ if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
+ return;
+ OpResult result = op->getResult(0);
+ if (!isa<VectorType>(result.getType()))
+ return;
+ xegpu::DistributeLayoutAttr resultLayout =
+ xegpu::getDistributeLayoutAttr(result);
+ if (!resultLayout)
+ return;
+
+ // Ensure that all feeding conversions have the same layout.
+ SmallVector<std::pair<OpOperand *, xegpu::ConvertLayoutOp>> conversions;
+ xegpu::DistributeLayoutAttr uniformConvSrcLayout;
+ for (OpOperand &operand : op->getOpOperands()) {
+ Value operandValue = operand.get();
+ if (!isa<VectorType>(operandValue.getType()))
+ continue;
+ auto operandConversionOp =
+ operandValue.getDefiningOp<xegpu::ConvertLayoutOp>();
+ // Those that are not a conversion should be easily materializable with
+ // another layout.
+ if (!operandConversionOp) {
+ Operation *definingOp = operandValue.getDefiningOp();
+ if (!definingOp || !xegpu::isTriviallyRematerializable(definingOp))
+ return;
+ continue;
+ }
+ xegpu::DistributeLayoutAttr input =
+ operandConversionOp.getEffectiveInputLayout();
+ if (!input)
+ return;
+ if (!uniformConvSrcLayout)
+ uniformConvSrcLayout = input;
+ else if (!uniformConvSrcLayout.isEqualTo(input))
+ return;
+ conversions.emplace_back(&operand, operandConversionOp);
+ }
+
+ if (!uniformConvSrcLayout)
+ return;
+
+ // Sink only when it leaves the op coarser.
+ // Conversion's source layout must be larger than the elemwise result
+ // layout.
+ SmallVector<int64_t> resultLayoutInstData =
+ resultLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> sourceInstData =
+ uniformConvSrcLayout.getEffectiveInstDataAsInt();
+ if (resultLayoutInstData.empty() || sourceInstData.empty() ||
+ computeProduct(sourceInstData) <= computeProduct(resultLayoutInstData))
+ return;
+
+ // Rewire the conversions, then rematerialize the remaining operands in the
+ // source layout.
+ llvm::SmallDenseSet<unsigned> rewiredOpIdxs;
+ for (auto [operand, convert] : conversions) {
+ operand->set(convert.getSource());
+ rewiredOpIdxs.insert(operand->getOperandNumber());
+ }
+ for (OpOperand &operand : op->getOpOperands()) {
+ Value operandValue = operand.get();
+ if (!isa<VectorType>(operandValue.getType()))
+ continue;
+ if (rewiredOpIdxs.contains(operand.getOperandNumber()))
+ continue;
+ Operation *definingOp = operandValue.getDefiningOp();
+ assert(definingOp && xegpu::isTriviallyRematerializable(definingOp) &&
+ "operand should have been rejected above");
+ // Rematerialize with uniform source layout.
+ builder.setInsertionPointAfter(definingOp);
+ Operation *clone = builder.clone(*definingOp);
+ OpResult cloneResult =
+ clone->getResult(cast<OpResult>(operandValue).getResultNumber());
+ xegpu::removeLayoutAttr(cloneResult);
+ xegpu::setDistributeLayoutAttr(cloneResult, uniformConvSrcLayout);
+ operand.set(cloneResult);
+ }
+
+ // Run the op in the source layout and bridge its result back, so that
+ // `getConsumerLayoutAt` now reports the source layout for every operand.
+ builder.setInsertionPointAfterValue(result);
+ auto newConvOp = xegpu::ConvertLayoutOp::create(
+ builder, op->getLoc(), result.getType(), result, uniformConvSrcLayout,
+ resultLayout);
+ result.replaceAllUsesExcept(newConvOp.getResult(), newConvOp);
+ xegpu::removeLayoutAttr(result);
+ xegpu::setDistributeLayoutAttr(result, uniformConvSrcLayout);
+
+ for (auto [operand, convert] : conversions)
+ if (convert.getResult().use_empty())
+ deadConverts.insert(convert);
+ });
+
+ // Memory effects block dce
+ for (xegpu::ConvertLayoutOp convert : deadConverts)
+ convert.erase();
+}
+
/// Update the function arguments and results with the layouts.
static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,
mlir::FunctionOpInterface funcOp,
@@ -1986,4 +2108,6 @@ void XeGPUPropagateLayoutPass::runOnOperation() {
signalPassFailure();
return;
}
+ if (layoutKind == xegpu::LayoutKind::InstData)
+ xegpu::sinkElementwiseConversions(builder, getOperation());
}
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 99e3f3be31c18..d822b487e5388 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -316,8 +316,6 @@ func.func @conflict_while_pass_through(%cond: i1) {
// CHECK: scf.yield %[[ADD]] : vector<16x16xf16>
// CHECK: }
// CHECK: %[[CVT:.*]] = xegpu.convert_layout %[[FOR]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
-// CHECK-SAME: : vector<16x16xf16>
// CHECK: %[[EXP:.*]] = math.exp %[[CVT]]
// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<16x16xf16>
// CHECK: return
@@ -423,25 +421,3 @@ func.func @extract_source_conflict_with_order() -> vector<16x32xf16> {
return %1 : vector<16x32xf16>
}
}
-
-// -----
-
-// Producer carries a layout that
diff ers from the input_layout declared on a
-// downstream xegpu.convert_layout consumer. ResolveLayoutConflicts must insert
-// a bridging convert_layout between the producer and the existing
-// convert_layout so the consumer's stated input_layout is honored.
-// CHECK-LABEL: func.func @convert_layout_bridge_input_mismatch
-// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
-// CHECK-NEXT: %[[BRIDGE:.*]] = xegpu.convert_layout %[[V0]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
-// CHECK-SAME: : vector<32x32xf16>
-gpu.module @test_convert_layout_bridge {
-func.func @convert_layout_bridge_input_mismatch() {
- %0 = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
- %1 = xegpu.convert_layout %0
- <{input_layout = #xegpu.layout<inst_data = [16, 16]>,
- target_layout = #xegpu.layout<inst_data = [32, 16]>}>
- : vector<32x32xf16>
- return
-}
-}
diff --git a/mlir/test/Dialect/XeGPU/sink-elementwise-conversions.mlir b/mlir/test/Dialect/XeGPU/sink-elementwise-conversions.mlir
new file mode 100644
index 0000000000000..4c26e6ef99c66
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/sink-elementwise-conversions.mlir
@@ -0,0 +1,120 @@
+// RUN: mlir-opt --test-xegpu-sink-elementwise-conversions -allow-unregistered-dialect --split-input-file %s | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: func.func @convert_layout_bridge_input_mismatch
+// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
+// CHECK-NEXT: %[[BRIDGE:.*]] = xegpu.convert_layout %[[V0]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [32, 16]>}>
+// CHECK-SAME: : vector<32x32xf16>
+gpu.module @test_convert_layout_bridge {
+func.func @convert_layout_bridge_input_mismatch() {
+ %0 = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
+ %1 = xegpu.convert_layout %0
+ <{input_layout = #xegpu.layout<inst_data = [16, 16]>,
+ target_layout = #xegpu.layout<inst_data = [32, 16]>}>
+ : vector<32x32xf16>
+ return
+}
+}
+
+// -----
+
+// CHECK-LABEL: func.func @sink_conversion_out_of_elementwise
+// CHECK-DAG: %[[MASK:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<16x64xi1>
+// CHECK-DAG: %[[DATA:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<16x64xf32>
+// CHECK-DAG: %[[PAD:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} dense<0xFF800000> : vector<16x64xf32>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[DATA]], %[[PAD]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<16x64xi1>, vector<16x64xf32>
+// CHECK: %[[CVT:.*]] = xegpu.convert_layout %[[SEL]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [1, 16]>}>
+// CHECK-SAME: : vector<16x64xf32>
+// CHECK: %{{.*}} = vector.multi_reduction <maximumf>, %[[CVT]], %{{.*}}
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test_sink_conversion {
+func.func @sink_conversion_out_of_elementwise() {
+ %mask = "some_op"() {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 16]>
+ } : () -> vector<16x64xi1>
+ %cmask = xegpu.convert_layout %mask <{input_layout = #xegpu.layout<inst_data = [8, 16]>,
+ target_layout = #xegpu.layout<inst_data = [1, 16]>}> : vector<16x64xi1>
+ %data = "some_op"() {
+ layout_result_0 = #xegpu.layout<inst_data = [8, 16]>
+ } : () -> vector<16x64xf32>
+ %cdata = xegpu.convert_layout %data <{input_layout = #xegpu.layout<inst_data = [8, 16]>,
+ target_layout = #xegpu.layout<inst_data = [1, 16]>}> : vector<16x64xf32>
+ %pad = arith.constant {
+ layout_result_0 = #xegpu.layout<inst_data = [1, 16]>
+ } dense<0xFF800000> : vector<16x64xf32>
+ %acc = arith.constant {
+ layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, dims = [1]>
+ } dense<0xFF800000> : vector<16xf32>
+
+ %sel = arith.select %cmask, %cdata, %pad {
+ layout_result_0 = #xegpu.layout<inst_data = [1, 16]>
+ } : vector<16x64xi1>, vector<16x64xf32>
+ %red = vector.multi_reduction <maximumf>, %sel, %acc {
+ layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, dims = [1]>
+ } [1] : vector<16x64xf32> to vector<16xf32>
+ return
+}
+}
+
+// -----
+
+// CHECK-LABEL: func.func @conflicting_operands_disagree_not_sunk
+// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
+// CHECK: %[[CVT0:.*]] = xegpu.convert_layout %[[V0]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [8, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+// CHECK: %[[V1:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [32, 16]>} : () -> vector<32x32xf16>
+// CHECK: %[[CVT1:.*]] = xegpu.convert_layout %[[V1]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+// CHECK: %{{.*}} = arith.addf %[[CVT0]], %[[CVT1]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : vector<32x32xf16>
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test_no_sink_on_disagreement {
+func.func @conflicting_operands_disagree_not_sunk() {
+ %0 = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : () -> vector<32x32xf16>
+ %c0 = xegpu.convert_layout %0
+ <{input_layout = #xegpu.layout<inst_data = [8, 16]>,
+ target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+ : vector<32x32xf16>
+ %1 = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [32, 16]>} : () -> vector<32x32xf16>
+ %c1 = xegpu.convert_layout %1
+ <{input_layout = #xegpu.layout<inst_data = [32, 16]>,
+ target_layout = #xegpu.layout<inst_data = [16, 16]>}>
+ : vector<32x32xf16>
+ %2 = arith.addf %c0, %c1 {layout_result_0 = #xegpu.layout<inst_data = [16, 16]>} : vector<32x32xf16>
+ return
+}
+}
+
+// -----
+
+// CHECK-LABEL: func.func @finer_operands_not_sunk
+// CHECK: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} : () -> vector<16x64xf32>
+// CHECK: %[[CVT0:.*]] = xegpu.convert_layout %[[V0]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [1, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+// CHECK: %[[V1:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} : () -> vector<16x64xf32>
+// CHECK: %[[CVT1:.*]] = xegpu.convert_layout %[[V1]]
+// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [1, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+// CHECK: %{{.*}} = arith.addf %[[CVT0]], %[[CVT1]]
+// CHECK-SAME: {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>} : vector<16x64xf32>
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test_no_sink_to_finer {
+func.func @finer_operands_not_sunk() {
+ %a = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} : () -> vector<16x64xf32>
+ %ca = xegpu.convert_layout %a
+ <{input_layout = #xegpu.layout<inst_data = [1, 16]>,
+ target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+ : vector<16x64xf32>
+ %b = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [1, 16]>} : () -> vector<16x64xf32>
+ %cb = xegpu.convert_layout %b
+ <{input_layout = #xegpu.layout<inst_data = [1, 16]>,
+ target_layout = #xegpu.layout<inst_data = [8, 16]>}>
+ : vector<16x64xf32>
+ %r = arith.addf %ca, %cb {layout_result_0 = #xegpu.layout<inst_data = [8, 16]>}
+ : vector<16x64xf32>
+ return
+}
+}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index a4813528e98dd..aacfd1a563d97 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -361,6 +361,36 @@ struct TestXeGPUResolveLayoutConflicts
}
};
+struct TestXeGPUSinkElementwiseConversions
+ : public PassWrapper<TestXeGPUSinkElementwiseConversions,
+ OperationPass<gpu::GPUModuleOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestXeGPUSinkElementwiseConversions)
+
+ StringRef getArgument() const final {
+ return "test-xegpu-sink-elementwise-conversions";
+ }
+
+ StringRef getDescription() const final {
+ return "Test sinking of XeGPU layout conversions past elementwise ops.";
+ }
+
+ void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
+ registry.insert<xegpu::XeGPUDialect>();
+ registry.insert<gpu::GPUDialect>();
+ registry.insert<vector::VectorDialect>();
+ }
+
+ TestXeGPUSinkElementwiseConversions() = default;
+ TestXeGPUSinkElementwiseConversions(
+ const TestXeGPUSinkElementwiseConversions &pass) = default;
+
+ void runOnOperation() override {
+ OpBuilder builder(&getContext());
+ xegpu::sinkElementwiseConversions(builder, getOperation());
+ }
+};
+
struct TestXeGPUArrayLengthOptimization
: public PassWrapper<TestXeGPUArrayLengthOptimization,
OperationPass<gpu::GPUModuleOp>> {
@@ -558,6 +588,7 @@ void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUSgToLaneDistribute>();
PassRegistration<TestXeGPUPropagateLayouts>();
PassRegistration<TestXeGPUResolveLayoutConflicts>();
+ PassRegistration<TestXeGPUSinkElementwiseConversions>();
PassRegistration<TestXeGPUArrayLengthOptimization>();
PassRegistration<TestXeGPUCoalesceGatherScatter>();
}
More information about the Mlir-commits
mailing list