[Mlir-commits] [mlir] [MLIR][XeGPU] Propagate layout onto loop-carried iter_arg entry edges (PR #198862)
Nishant Patel
llvmlistbot at llvm.org
Wed May 27 08:19:30 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/198862
>From 38be227ec4e224ef47d19474ab32127a469e3366 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 19 May 2026 16:12:47 +0000
Subject: [PATCH 1/4] [mlir][xegpu] Propagate layout onto loop-carried iter_arg
entry edges
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 45 +++++++++++++++++++
.../XeGPU/propagate-layout-subgroup.mlir | 2 +-
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 4 +-
3 files changed, 48 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ffcda32e55c04..cc6f8bba737f2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1734,6 +1734,47 @@ updateControlFlowOps(mlir::OpBuilder &builder,
return success();
}
+/// Parent-rooted counterpart of `updateControlFlowOps`. For each entry edge
+/// of a `RegionBranchOpInterface` op (i.e. the edges from the parent op into
+/// its regions on first execution), plant `layout_operand_N` on the parent
+/// op's entry operand whenever the tied region-entry input (typically a body
+/// block argument) has a propagated layout.
+///
+/// This is required because `getDistributeLayoutAttr(BlockArgument)`
+/// redirects to `getTemporaryLayout(tiedInit)` on the parent op. Without
+/// writing the entry-operand layout back here, that lookup returns null even
+/// when backward propagation has fully determined the layout (e.g. from a
+/// downstream `xegpu.dpas` anchor through a `scf.yield` back-edge), and
+/// downstream vector consumers are treated as having a layout-less producer.
+///
+/// Written against `RegionBranchOpInterface` so it covers loops, `scf.if`,
+/// `scf.while`, and any other op implementing the interface symmetrically.
+static LogicalResult
+updateRegionBranchEntryEdges(mlir::OpBuilder &builder,
+ mlir::RegionBranchOpInterface branchOp,
+ GetLayoutFnTy getLayoutOfValue) {
+ SmallVector<RegionSuccessor> entrySuccessors;
+ branchOp.getSuccessorRegions(RegionBranchPoint::parent(), entrySuccessors);
+ for (const RegionSuccessor &successor : entrySuccessors) {
+ if (successor.isParent())
+ continue;
+ OperandRange entryOperands = branchOp.getEntrySuccessorOperands(successor);
+ ValueRange successorInputs = branchOp.getSuccessorInputs(successor);
+ if (entryOperands.empty() || entryOperands.size() != successorInputs.size())
+ continue;
+ unsigned beginIdx = entryOperands.getBeginOperandIndex();
+ for (auto [i, successorInput] : llvm::enumerate(successorInputs)) {
+ if (!isa<VectorType, xegpu::TensorDescType>(successorInput.getType()))
+ continue;
+ xegpu::DistributeLayoutAttr layout = getLayoutOfValue(successorInput);
+ if (!layout)
+ continue;
+ xegpu::setTemporaryLayout(branchOp->getOpOperand(beginIdx + i), layout);
+ }
+ }
+ return success();
+}
+
/// Update the function arguments and results with the layouts.
static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,
mlir::FunctionOpInterface funcOp,
@@ -1828,6 +1869,10 @@ LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
r = updateControlFlowOps(builder, branchTermOp,
getXeGPULayoutForValue);
})
+ .Case([&](mlir::RegionBranchOpInterface branchOp) {
+ r = updateRegionBranchEntryEdges(builder, branchOp,
+ getXeGPULayoutForValue);
+ })
.Case([&](mlir::FunctionOpInterface funcOp) {
r = updateFunctionOpInterface(builder, funcOp,
getXeGPULayoutForValue);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index b4a34fb01456f..50128843474ce 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -260,7 +260,7 @@ gpu.module @test {
// CHECK-SAME: : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
// CHECK-NEXT: scf.yield %{{.*}} : vector<128x128xf32>
- // CHECK-NEXT: } {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}
+ // CHECK-NEXT: } {layout_operand_3 = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>, layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}
// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}>
%2 = scf.for %arg3 = %c0 to %c8192 step %c128 iter_args(%arg4 = %cst) -> (vector<128x128xf32>) {
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 3ceddde54dede..1bfa3b5b6d584 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -389,7 +389,7 @@ gpu.module @test {
// CHECK-NEXT: %[[T6:.*]] = xegpu.dpas %[[T4]], %[[T5]], %[[ARG6]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
// CHECK-NEXT: scf.yield %[[T6]] : vector<8x16xf32>
-// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+// CHECK-NEXT: } {layout_operand_3 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
// CHECK-NEXT: %[[T3:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK-NEXT: xegpu.store_nd %[[T2]], %[[T3]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
func.func @for_op(%arg0: memref<8x128xf16>, %arg1: memref<128x16xf16>, %arg2: memref<8x16xf32>) {
@@ -556,7 +556,7 @@ gpu.module @test {
// CHECK-NEXT: } do {
// CHECK-NEXT: ^bb0(%{{.*}}: vector<16xf32>, %{{.*}}: i32):
// CHECK: scf.yield {{.*}} : vector<16xf32>, i32
-// CHECK-NEXT: } attributes {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
+// CHECK-NEXT: } attributes {layout_operand_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>, layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
func.func @scf_while_and_condition(%arg0: memref<256xf32>, %arg1: memref<256xf32>) {
%c0 = arith.constant 0 : i32
%c16 = arith.constant 16 : i32
>From 5ab1ab5450513e8506669dd7acf3738b34bddf68 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Thu, 21 May 2026 20:45:10 +0000
Subject: [PATCH 2/4] Fix comment
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index cc6f8bba737f2..2448bc86f392c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1744,8 +1744,8 @@ updateControlFlowOps(mlir::OpBuilder &builder,
/// redirects to `getTemporaryLayout(tiedInit)` on the parent op. Without
/// writing the entry-operand layout back here, that lookup returns null even
/// when backward propagation has fully determined the layout (e.g. from a
-/// downstream `xegpu.dpas` anchor through a `scf.yield` back-edge), and
-/// downstream vector consumers are treated as having a layout-less producer.
+/// downward `xegpu.dpas` anchor through a `scf.yield` back-edge), and
+/// vector consumers are treated as having a layout-less producer.
///
/// Written against `RegionBranchOpInterface` so it covers loops, `scf.if`,
/// `scf.while`, and any other op implementing the interface symmetrically.
>From 5338a076aabdb49523a81de111da6bfccd7af618 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 26 May 2026 20:40:00 +0000
Subject: [PATCH 3/4] Address feedback
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 14 ++++++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 50 ++++++++++++-------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 45 +----------------
3 files changed, 48 insertions(+), 61 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 8ffba7ff36208..8e679522ee447 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -14,6 +14,8 @@
#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
namespace mlir {
@@ -39,6 +41,18 @@ LogicalResult propagateLayouts(OpBuilder &builder, Operation *target,
LogicalResult resolveLayoutConflicts(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)>;
+
+/// Propagate layouts from a region branch op's region entry block arguments
+/// back to its init operands. The block argument's layout is obtained via
+/// `getLayoutOfValue`; the matching layout is then recorded on each init
+/// operand that flows into that block argument (e.g. scf.for's iter_args
+/// inits), and on tensor descriptor block argument types.
+LogicalResult propagateRegionArgsToInits(RegionBranchOpInterface regionOp,
+ GetLayoutFnTy getLayoutOfValue);
+
/// Attach layout attributes to all vector-type operands of operations within
/// the given operation's nested region. Reports an error if any vector operand
/// lacks a layout attribute.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 813e9608dbf8e..b651bc5ebbc88 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -216,9 +216,11 @@ static void propagateRegionResultsToYieldOperands(
// Propagate layout from region arguments to region op's init operands. This
// sets the temporary layout for region arguments and init operands.
-static void propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp) {
+LogicalResult
+xegpu::propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp,
+ xegpu::GetLayoutFnTy getLayoutOfValue) {
// Iterate all regions of the region op. For each block argument that has a
- // layout (determined from its use points), trace back to find the
+ // layout (obtained via `getLayoutOfValue`), trace back to find the
// corresponding init operand of the regionOp and set the layout on it.
// This works generically for scf.for, scf.while, and other
// RegionBranchOpInterface ops.
@@ -229,7 +231,7 @@ static void propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp) {
// the induction variable is a block arg but not a successor input.
ValueRange successorInputs = regionOp.getSuccessorInputs(regionSuccessor);
for (auto [inputIdx, regionArg] : llvm::enumerate(successorInputs)) {
- auto layout = getLayoutFromUsePoints(regionArg);
+ auto layout = getLayoutOfValue(regionArg);
if (!layout)
continue;
@@ -250,6 +252,7 @@ static void propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp) {
}
}
}
+ return success();
}
// Prerequisite for Layout Recovery
@@ -285,7 +288,8 @@ bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
auto processFunc = [&](Region &body, StringRef funcName) {
walkRegionBackward(body, [&](Operation *op) {
if (auto regionOp = dyn_cast<mlir::RegionBranchOpInterface>(op)) {
- propagateRegionArgsToInits(regionOp);
+ (void)xegpu::propagateRegionArgsToInits(regionOp,
+ getLayoutFromUsePoints);
} else if (auto yieldOp =
dyn_cast<mlir::RegionBranchTerminatorOpInterface>(op)) {
propagateRegionResultsToYieldOperands(yieldOp);
@@ -1009,43 +1013,47 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"laneData must be available for all dimensions");
- size_t dim = srcShape.size() - 1;
+ size_t innerMostDim = srcShape.size() - 1;
int64_t sgDataValue = -1;
int64_t instDataValue = -1;
int64_t laneDataValue = -1;
- const int subgroupSize = uArch->getSubgroupSize();
if (srcElemTyBitWidth > resElemTyBitWidth) {
// When casting to a smaller bitwidth, multiply the result layout
// accordingly to ensure it can be divided by the ratio back to the
// source layout.
int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
- int innermostDimLaneLayout = subgroupSize;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[dim];
- while ((sgDataValue <= resShape[dim]) &&
+ sgDataValue = sgData[innerMostDim];
+ while ((sgDataValue <= resShape[innerMostDim]) &&
(sgDataValue % bitWidthRatio) != 0)
sgDataValue *= 2;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[dim];
+ instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
// Adjust instDataValue so it still fits within an instruction after
// dividing by bitWidthRatio
- while ((instDataValue <= resShape[dim]) &&
+ while ((instDataValue <= resShape[innerMostDim]) &&
(instDataValue % (innermostDimLaneLayout * bitWidthRatio) != 0))
instDataValue *= 2;
- assert((resShape[dim] % instDataValue) == 0 &&
+ assert((resShape[innerMostDim] % instDataValue) == 0 &&
"resShape, instData, and lanelayout for innermost must be 2^n !");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[dim];
- while ((laneDataValue <= resShape[dim]) &&
+ laneDataValue = laneData[innerMostDim];
+ while ((laneDataValue <= resShape[innerMostDim]) &&
(laneDataValue % bitWidthRatio != 0))
laneDataValue *= 2;
}
// Now set only instData and laneData, preserving sgData
xegpu::DistributeLayoutAttr resLayout;
- resLayout = consumerLayout.setDimData(dim, sgDataValue, instDataValue,
- laneDataValue);
+ resLayout = consumerLayout.setDimData(innerMostDim, sgDataValue,
+ instDataValue, laneDataValue);
return resLayout;
}
return consumerLayout;
@@ -1074,6 +1082,8 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"consumer layout rank must match source shape rank");
@@ -1084,7 +1094,6 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
// Interleave doubles the innermost dimension (ratio = 2)
constexpr int ratio = 2;
- int innermostDimLaneLayout = uArch->getSubgroupSize();
if (layoutKind == xegpu::LayoutKind::Subgroup) {
sgDataValue = sgData[innerMostDim];
@@ -1094,6 +1103,9 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
sgDataValue *= ratio;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
// Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
// ratio) when inferring the source layout
while ((instDataValue <= srcShape[innerMostDim]) &&
@@ -1456,6 +1468,8 @@ getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
int kDimSize = subgroupSize;
if (isDpasMx) {
auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (supportedKLen.empty())
+ return std::nullopt;
kDimSize = supportedKLen[0];
}
@@ -1910,7 +1924,7 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For non-anchor ops, derive the operand layout from the op's result
// layout via op-specific semantics.
xegpu::DistributeLayoutAttr resLayout;
- if (op->getNumResults() == 1)
+ if (op->getNumResults() == 1 || isa<vector::DeinterleaveOp>(op))
resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
return inferSourceLayoutFromResultForNonAnchorOp(operand, resLayout);
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 2448bc86f392c..b8efe38be75bb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1734,47 +1734,6 @@ updateControlFlowOps(mlir::OpBuilder &builder,
return success();
}
-/// Parent-rooted counterpart of `updateControlFlowOps`. For each entry edge
-/// of a `RegionBranchOpInterface` op (i.e. the edges from the parent op into
-/// its regions on first execution), plant `layout_operand_N` on the parent
-/// op's entry operand whenever the tied region-entry input (typically a body
-/// block argument) has a propagated layout.
-///
-/// This is required because `getDistributeLayoutAttr(BlockArgument)`
-/// redirects to `getTemporaryLayout(tiedInit)` on the parent op. Without
-/// writing the entry-operand layout back here, that lookup returns null even
-/// when backward propagation has fully determined the layout (e.g. from a
-/// downward `xegpu.dpas` anchor through a `scf.yield` back-edge), and
-/// vector consumers are treated as having a layout-less producer.
-///
-/// Written against `RegionBranchOpInterface` so it covers loops, `scf.if`,
-/// `scf.while`, and any other op implementing the interface symmetrically.
-static LogicalResult
-updateRegionBranchEntryEdges(mlir::OpBuilder &builder,
- mlir::RegionBranchOpInterface branchOp,
- GetLayoutFnTy getLayoutOfValue) {
- SmallVector<RegionSuccessor> entrySuccessors;
- branchOp.getSuccessorRegions(RegionBranchPoint::parent(), entrySuccessors);
- for (const RegionSuccessor &successor : entrySuccessors) {
- if (successor.isParent())
- continue;
- OperandRange entryOperands = branchOp.getEntrySuccessorOperands(successor);
- ValueRange successorInputs = branchOp.getSuccessorInputs(successor);
- if (entryOperands.empty() || entryOperands.size() != successorInputs.size())
- continue;
- unsigned beginIdx = entryOperands.getBeginOperandIndex();
- for (auto [i, successorInput] : llvm::enumerate(successorInputs)) {
- if (!isa<VectorType, xegpu::TensorDescType>(successorInput.getType()))
- continue;
- xegpu::DistributeLayoutAttr layout = getLayoutOfValue(successorInput);
- if (!layout)
- continue;
- xegpu::setTemporaryLayout(branchOp->getOpOperand(beginIdx + i), layout);
- }
- }
- return success();
-}
-
/// Update the function arguments and results with the layouts.
static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,
mlir::FunctionOpInterface funcOp,
@@ -1870,8 +1829,8 @@ LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
getXeGPULayoutForValue);
})
.Case([&](mlir::RegionBranchOpInterface branchOp) {
- r = updateRegionBranchEntryEdges(builder, branchOp,
- getXeGPULayoutForValue);
+ r = xegpu::propagateRegionArgsToInits(branchOp,
+ getXeGPULayoutForValue);
})
.Case([&](mlir::FunctionOpInterface funcOp) {
r = updateFunctionOpInterface(builder, funcOp,
>From 2fc3832280bcc6a6ef11953f4586b2ed532e2ed2 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 27 May 2026 15:18:52 +0000
Subject: [PATCH 4/4] Change function name
---
.../Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 11 ++++++-----
1 file changed, 6 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 3debc324fdf8f..4bde8a40e8c91 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1803,7 +1803,8 @@ LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
return success();
}
// Helper to convert LayoutInfo to xegpu::LayoutAttr.
- auto getXeGPULayoutForValue = [&](Value val) -> xegpu::DistributeLayoutAttr {
+ auto getLayoutFromPropagation =
+ [&](Value val) -> xegpu::DistributeLayoutAttr {
LayoutInfo layout = analysis.getLayoutInfo(val);
if (auto opResult = dyn_cast<OpResult>(val)) {
Operation *defOp = opResult.getDefiningOp();
@@ -1834,18 +1835,18 @@ LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
TypeSwitch<Operation *>(&op)
.Case([&](mlir::RegionBranchTerminatorOpInterface branchTermOp) {
r = updateControlFlowOps(builder, branchTermOp,
- getXeGPULayoutForValue);
+ getLayoutFromPropagation);
})
.Case([&](mlir::RegionBranchOpInterface branchOp) {
r = xegpu::propagateRegionArgsToInits(branchOp,
- getXeGPULayoutForValue);
+ getLayoutFromPropagation);
})
.Case([&](mlir::FunctionOpInterface funcOp) {
r = updateFunctionOpInterface(builder, funcOp,
- getXeGPULayoutForValue);
+ getLayoutFromPropagation);
})
.Default([&](Operation *op) {
- r = updateOp(builder, op, getXeGPULayoutForValue);
+ r = updateOp(builder, op, getLayoutFromPropagation);
});
if (failed(r)) {
op.emitError("Failed to update operation with the layout.");
More information about the Mlir-commits
mailing list