[Mlir-commits] [mlir] [MLIR][XeGPU] Add local forward layout propagation (PR #208932)
Jianhui Li
llvmlistbot at llvm.org
Tue Jul 14 19:21:36 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/208932
>From d83861877d11a1f6e2d57b933be7cbc1a769269b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 11 Jul 2026 18:16:28 +0000
Subject: [PATCH] [MLIR][XeGPU] Add local forward layout propagation
Backward layout propagation only assigns layouts to values that are
(transitively) consumed by an anchor op. A value whose only consumer is,
e.g., the next iteration of a loop is left without a layout.
Add a local forward-fill step, run after the backward materialization walk
in propagateLayouts(): it visits ops in producer-first order and, for any
un-laid-out vector result, infers the layout from the op's already-known
operand layouts via a new inferResultLayoutFromSourceForNonAnchorOp
dispatcher (covering elementwise, transpose, and shape_cast; other ops are
left as TODO), then stamps it with setDistributeLayoutAttr.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 43 ++++++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 123 ++++++++++++++++++
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 56 ++++++++
.../XeGPU/propagate-layout-subgroup.mlir | 35 +++++
4 files changed, 257 insertions(+)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 75fcb156a9d0d..040227996f0ae 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -83,6 +83,15 @@ dropSgLayoutAndDataOnAttrs(ArrayRef<NamedAttribute> attrs);
/// any DistributeLayoutAttr found.
SmallVector<NamedAttribute> dropInstDataOnAttrs(ArrayRef<NamedAttribute> attrs);
+//===----------------------------------------------------------------------===//
+// Backward layout inference (result layout -> source layout)
+//===----------------------------------------------------------------------===//
+//
+// The infer*SourceLayout helpers below derive the layout of an operation's
+// source operand from the layout of its result. They implement the per-op
+// transfer functions used by the backward layout propagation analysis, which
+// flows layouts from anchor ops (dpas, store_nd, ...) back to their producers.
+
/// Infers the source layout attribute for a broadcast operation given the
/// result layout attribute, result shape, and source shape.
DistributeLayoutAttr inferBroadcastSourceLayout(DistributeLayoutAttr resLayout,
@@ -159,6 +168,40 @@ DistributeLayoutAttr
inferSourceLayoutFromResultForNonAnchorOp(OpOperand &operand,
DistributeLayoutAttr resLayout);
+//===----------------------------------------------------------------------===//
+// Forward layout inference (source layout -> result layout)
+//===----------------------------------------------------------------------===//
+//
+// The infer*ResultLayout helpers below are the forward counterparts of the
+// infer*SourceLayout helpers above: given the layout of an operation's source
+// operand they derive the layout of its result. They are used by the local
+// forward-fill step in XeGPUPropagateLayout that assigns layouts to values not
+// reached by the backward propagation analysis (e.g. loop-carried values whose
+// only consumer is the next iteration).
+
+/// Infers the result layout attribute for a transpose operation given the
+/// source layout attribute and permutation. Inverse of
+/// inferTransposeSourceLayout.
+DistributeLayoutAttr inferTransposeResultLayout(DistributeLayoutAttr srcLayout,
+ ArrayRef<int64_t> permutation);
+
+/// Infers the result layout attribute for a shape cast operation given the
+/// source layout attribute, source shape, and result shape. Inverse of
+/// inferShapeCastSourceLayout. Returns nullptr for shape-cast patterns whose
+/// forward direction is ambiguous (e.g. unit-dim expansion).
+DistributeLayoutAttr inferShapeCastResultLayout(DistributeLayoutAttr srcLayout,
+ ArrayRef<int64_t> srcShape,
+ ArrayRef<int64_t> resShape);
+
+/// Infers the result layout attribute for a non-anchor operation from the
+/// layouts of its source operands (the forward counterpart of
+/// inferSourceLayoutFromResultForNonAnchorOp). `operandLayouts` is indexed by
+/// operand number; entries may be null for operands without a known layout.
+/// Returns nullptr when no forward rule applies (the result is then left
+/// un-laid-out).
+DistributeLayoutAttr inferResultLayoutFromSourceForNonAnchorOp(
+ Operation *op, ArrayRef<DistributeLayoutAttr> operandLayouts);
+
/// Note on the `consumerLayout` argument used by the consumer-driven setup* /
/// complete* helpers below:
///
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 8da151a333fb3..d68f81c7fd3cd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -835,6 +835,129 @@ xegpu::inferShapeCastSourceLayout(xegpu::DistributeLayoutAttr resLayout,
return nullptr;
}
+//===----------------------------------------------------------------------===//
+// Forward layout inference (source layout -> result layout)
+//===----------------------------------------------------------------------===//
+
+/// Infers the result layout attribute for a transpose operation given the
+/// source layout attribute and permutation.
+///
+/// vector.transpose semantics is `result[i] = source[permutation[i]]`, so
+/// `result_layout[i] = source_layout[permutation[i]]`, which is exactly
+/// `srcLayout.transposeDims(permutation)`. This is the inverse of
+/// inferTransposeSourceLayout (which applies the inverse permutation).
+xegpu::DistributeLayoutAttr
+xegpu::inferTransposeResultLayout(xegpu::DistributeLayoutAttr srcLayout,
+ ArrayRef<int64_t> permutation) {
+ return srcLayout.transposeDims(permutation);
+}
+
+/// Infers the result layout attribute for a shape cast operation given the
+/// source layout attribute, source shape, and result shape. This is the
+/// inverse of inferShapeCastSourceLayout: a dim-split (src -> res) is undone by
+/// collapsing the split groups, and a dim-collapse (src -> res) is undone by
+/// expanding the collapsed groups. The unit-dim-expansion case is not inverted
+/// here because recovering which result dims are the expanded unit dims would
+/// require the SliceAttr the backward direction produces; such patterns return
+/// nullptr (leaving the result un-laid-out).
+xegpu::DistributeLayoutAttr
+xegpu::inferShapeCastResultLayout(xegpu::DistributeLayoutAttr srcLayout,
+ ArrayRef<int64_t> srcShape,
+ ArrayRef<int64_t> resShape) {
+ // Case: source dims were split into result dims (forward of use case 2 in
+ // inferShapeCastSourceLayout). Undo by expanding each source dim into its
+ // group of result dims.
+ SmallVector<SmallVector<int64_t>> splitDimGroups;
+ if (xegpu::matchSplitDimExpansion(srcShape, resShape, splitDimGroups)) {
+ auto resLayout = srcLayout;
+ // Process source dims from innermost to outermost so that expanding a dim
+ // does not shift the indices of dims not yet processed.
+ for (int64_t srcIdx = static_cast<int64_t>(splitDimGroups.size()) - 1;
+ srcIdx >= 0; --srcIdx) {
+ ArrayRef<int64_t> resDims = splitDimGroups[srcIdx];
+ if (resDims.size() <= 1)
+ continue;
+ SmallVector<int64_t> targetShape;
+ targetShape.reserve(resDims.size());
+ for (int64_t d : resDims)
+ targetShape.push_back(resShape[d]);
+ resLayout = resLayout.expandDim(srcIdx, targetShape);
+ }
+ return resLayout;
+ }
+
+ // Case: source dims were collapsed into result dims (forward of use case 3).
+ // Undo by collapsing each group of source dims into its single result dim.
+ SmallVector<SmallVector<int64_t>> collapseDims;
+ if (xegpu::matchDimCollapse(srcShape, resShape, collapseDims)) {
+ auto resLayout = srcLayout;
+ // Process result dims from innermost to outermost so that collapsing a
+ // group does not shift the indices of groups not yet processed.
+ for (int64_t dstIdx = static_cast<int64_t>(collapseDims.size()) - 1;
+ dstIdx >= 0; --dstIdx) {
+ ArrayRef<int64_t> srcDims = collapseDims[dstIdx];
+ // A result dim with no backing source dims is a trailing/leading unit
+ // dim; its forward inference is ambiguous, so bail out.
+ if (srcDims.empty())
+ return nullptr;
+ if (srcDims.size() == 1)
+ continue;
+ resLayout = resLayout.collapseDims(llvm::to_vector(srcDims));
+ }
+ return resLayout;
+ }
+
+ return nullptr;
+}
+
+/// Infers the result layout attribute for a non-anchor operation from the
+/// layouts of its source operands. Forward counterpart of
+/// inferSourceLayoutFromResultForNonAnchorOp.
+xegpu::DistributeLayoutAttr xegpu::inferResultLayoutFromSourceForNonAnchorOp(
+ Operation *op, ArrayRef<xegpu::DistributeLayoutAttr> operandLayouts) {
+ if (op->getNumResults() != 1)
+ return nullptr;
+
+ // For vector::TransposeOp, infer the result layout from the source layout.
+ if (auto transpose = dyn_cast<vector::TransposeOp>(op)) {
+ if (!operandLayouts[0])
+ return nullptr;
+ return xegpu::inferTransposeResultLayout(operandLayouts[0],
+ transpose.getPermutation());
+ }
+
+ // For vector::ShapeCastOp, infer the result layout from the source layout.
+ if (auto shapeCast = dyn_cast<vector::ShapeCastOp>(op)) {
+ if (!operandLayouts[0])
+ return nullptr;
+ return xegpu::inferShapeCastResultLayout(
+ operandLayouts[0], shapeCast.getSourceVectorType().getShape(),
+ shapeCast.getResultVectorType().getShape());
+ }
+
+ // For elementwise operations, all operands and the result share the same
+ // layout. Use the first operand that carries a layout.
+ if (OpTrait::hasElementwiseMappableTraits(op)) {
+ for (xegpu::DistributeLayoutAttr layout : operandLayouts)
+ if (layout)
+ return layout;
+ return nullptr;
+ }
+
+ // TODO: add forward inference rules for the remaining ops; their result is
+ // left un-laid-out until then.
+ // - vector::BroadcastOp: the forward direction is under-determined. The
+ // backward rule (inferBroadcastSourceLayout) either sets broadcast dims to
+ // unit data (losing the original data on those dims) or wraps the result
+ // in a SliceAttr; neither is generally invertible from the source layout
+ // alone, so a forward rule must decide how to distribute the new/stretched
+ // dims.
+ // - vector::BitCastOp, vector::MultiDimReductionOp / vector::ReductionOp,
+ // vector::InterleaveOp / vector::DeinterleaveOp, and the insert / extract
+ // / strided-slice family.
+ return nullptr;
+}
+
/// Infers the layout attribute for mask and offset operand for Chunked load
/// and store, given the anchor layout attribute for the value being load/store.
xegpu::DistributeLayoutAttr xegpu::inferMaskOffsetLayoutForScatterIO(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 1b1abcc4b0e4f..68f8abd1ed92d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1695,6 +1695,57 @@ ResolveLayoutConflicts::resolveTensorDescConsumer(OpOperand &operand) {
}
using GetLayoutFnTy = function_ref<xegpu::DistributeLayoutAttr(Value)>;
+
+/// Local forward layout propagation, run after the backward propagation
+/// analysis has materialized its layouts. Backward propagation only assigns
+/// layouts to values that are (transitively) consumed by an anchor op; a value
+/// whose only consumer is, e.g., the next iteration of a loop (its producing op
+/// is not on any anchor's backward slice) is left without a layout. This walk
+/// visits ops in producer-first (forward) order and, for any vector result that
+/// still lacks a layout, infers it from the already-known layouts of the op's
+/// operands via inferResultLayoutFromSourceForNonAnchorOp, then stamps it with
+/// setDistributeLayoutAttr.
+///
+/// A single forward pass suffices: forward-only values form producer ->
+/// consumer chains that the pre-order walk visits in dependency order, so an
+/// operand's layout is already assigned by the time its consumer is visited.
+static void forwardFillLayouts(Operation *root) {
+ root->walk([&](Operation *op) {
+ // Anchor ops carry authoritative layouts; region ops and their terminators
+ // are handled by the control-flow propagation; function ops carry no result
+ // layouts here.
+ if (isa<xegpu::AnchorLayoutInterface, RegionBranchOpInterface,
+ RegionBranchTerminatorOpInterface, FunctionOpInterface>(op))
+ return;
+ if (op->getNumResults() != 1)
+ return;
+ OpResult result = op->getResult(0);
+ if (!isa<VectorType>(result.getType()))
+ return;
+ // Skip results that already have a layout (from backward propagation or an
+ // earlier forward-fill step).
+ if (xegpu::getDistributeLayoutAttr(result))
+ return;
+
+ // Gather operand layouts, indexed by operand number.
+ SmallVector<xegpu::DistributeLayoutAttr> operandLayouts;
+ operandLayouts.reserve(op->getNumOperands());
+ bool anyAssigned = false;
+ for (Value operand : op->getOperands()) {
+ auto layout = xegpu::getDistributeLayoutAttr(operand);
+ operandLayouts.push_back(layout);
+ anyAssigned |= (layout != nullptr);
+ }
+ if (!anyAssigned)
+ return;
+
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::inferResultLayoutFromSourceForNonAnchorOp(op, operandLayouts);
+ if (layout)
+ xegpu::setDistributeLayoutAttr(result, layout);
+ });
+}
+
/// Update an operation with the layout of its results. If the result type is
/// a vector type, a temporary layout attribute is added to the operation. If
/// the result type is a tensor descriptor type, the type is updated with the
@@ -1918,6 +1969,11 @@ LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
if (walkResult.wasInterrupted())
return failure();
+ // Backward propagation only reaches values consumed by anchor ops. Run a
+ // local forward pass to fill in layouts for values it could not reach (e.g.
+ // loop-carried values whose only consumer is the next iteration).
+ forwardFillLayouts(op);
+
return success();
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index fa00c1d894d8f..5c29695bafa31 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -672,3 +672,38 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+// Forward layout fill: %sum is a loop-carried value whose only consumer is the
+// next iteration (via iter_arg %acc), so backward propagation from anchor ops
+// never reaches it. The local forward pass derives its layout from the operand
+// %acc (elementwise rule), and the scf.for's second result / yield inherit it.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @forward_fill_loop_carried(
+gpu.func @forward_fill_loop_carried(%arg0: memref<128x64xf16>, %arg1: memref<64x64xf16>, %arg2: memref<128x64xf32>) kernel attributes {known_block_size = array<i32: 128, 1, 1>} {
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c128 = arith.constant 128 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<128x64xf32>
+ %cst_0 = arith.constant dense<0.000000e+00> : vector<128xf32>
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<128x64xf16> -> !xegpu.tensor_desc<128x64xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>}> : !xegpu.tensor_desc<128x64xf16> -> vector<128x64xf16>
+ %2 = xegpu.create_nd_tdesc %arg1 : memref<64x64xf16> -> !xegpu.tensor_desc<64x64xf16>
+ %3 = xegpu.load_nd %2[%c0, %c0] <{layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [64, 64]>}> : !xegpu.tensor_desc<64x64xf16> -> vector<64x64xf16>
+ // CHECK: scf.for
+ %4:2 = scf.for %arg4 = %c0 to %c128 step %c64 iter_args(%acc0 = %cst, %acc = %cst_0) -> (vector<128x64xf32>, vector<128xf32>) {
+ %8 = xegpu.dpas %1, %3, %acc0 {layout_a = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>, layout_b = #xegpu.layout<sg_layout = [1, 1], sg_data = [64, 64]>, layout_cd = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>} : vector<128x64xf16>, vector<64x64xf16>, vector<128x64xf32> -> vector<128x64xf32>
+ %9 = vector.broadcast %acc : vector<128xf32> to vector<64x128xf32>
+ %10 = vector.transpose %9, [1, 0] : vector<64x128xf32> to vector<128x64xf32>
+ %11 = arith.mulf %8, %10 : vector<128x64xf32>
+ // The forward pass assigns %sum a layout inherited from %acc.
+ // CHECK: arith.addf
+ // CHECK-SAME: layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 8], sg_data = [64, 16], order = [0, 1]>, dims = [0]>
+ %sum = arith.addf %acc, %acc : vector<128xf32>
+ scf.yield %11, %sum : vector<128x64xf32>, vector<128xf32>
+ }
+ %5 = xegpu.create_nd_tdesc %arg2 : memref<128x64xf32> -> !xegpu.tensor_desc<128x64xf32>
+ xegpu.store_nd %4#0, %5[%c0, %c0] <{layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>}> : vector<128x64xf32>, !xegpu.tensor_desc<128x64xf32>
+ gpu.return
+}
+}
More information about the Mlir-commits
mailing list