[Mlir-commits] [mlir] [mlir][xegpu] SLM Privatization for WG-to-SG Distribution (PR #208921)
Jianhui Li
llvmlistbot at llvm.org
Sat Jul 11 08:31:35 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/208921
Adds a pre-phase to xegpu-wg-to-sg-distribute that demotes shared local memory (space 3) scratch buffers to subgroup-private memory (space 4, register-backed) when they're accessed identically by every subgroup.
When a buffer qualifies: all load_matrix/store_matrix on it share the same sg_layout, sg_data, data shape, and offsets; the
workgroup tile distributes evenly across subgroups (no broadcast/overlap); and the source alloca has exactly one create_mem_desc view (doesn't escape). It flips the memref to space 4, shrinks the buffer + mem_desc to the per-subgroup size, and re-indexes matrix ops with local, subgroup-id-free offsets.
>From c589918d9f4111e049d6165a56fb7fc2bfa6d344 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 8 Jul 2026 23:30:49 +0000
Subject: [PATCH 1/9] [MLIR][XeGPU] Prefer slice layout in LayoutInfo::meet
During backward layout propagation a value can receive different
candidate layouts from its consumers. Previously LayoutInfo::meet kept
whichever layout was assigned first, so a value feeding both a
slice-producing consumer (e.g. a broadcast whose inferred source layout
is a slice) and a plain-layout consumer (e.g. a dpas_mx scale operand)
could end up with the plain layout depending on visitation order.
meet() now prefers the slice layout when both sides are assigned
(falling back to lhs when both/neither are slices). For that preference
to take effect, operator== compares the stored layout (via isEqualTo)
rather than only the assigned bit; otherwise the dataflow framework
treats the refined result as a no-op and discards it. The refinement
only moves a plain layout toward a slice layout, so it stays monotonic.
Adds a propagate-layout-inst-data test and updates the a-operand
lane_data in the simple_mxfp_gemm integration test to match the HW
constraint.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 25 +++++++++++-----
.../XeGPU/propagate-layout-inst-data.mlir | 29 +++++++++++++++++++
.../Dialect/XeGPU/WG/simple_mxfp_gemm.mlir | 2 +-
3 files changed, 47 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 64d0d8063b7ff..ff9239b637a10 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -71,12 +71,13 @@ namespace {
/// Given this, LayoutInfo satisifies the following properties:
/// 1) A LayoutInfo value can be in one of two states - `assigned` or `not
/// assigned`.
-/// 2) Two LayoutInfo values are equal if they are both assigned or
-/// both not assigned. The concrete value of assigned state does not matter.
+/// 2) Two LayoutInfo values are equal if they are both not assigned, or both
+/// assigned with the same layout.
/// 3) The meet operator works as follows:
-/// - If current state is assigned, return the current state. (already
-/// a unique layout is assigned. don't change it)
-/// - Otherwise, return the other state.
+/// - If only one side is assigned, return that side.
+/// - If both sides are assigned, prefer the side carrying a slice layout.
+/// If both (or neither) are slice layouts, prefer the lhs (current
+/// state) so an already assigned unique layout is not changed.
struct LayoutInfo {
private:
@@ -86,10 +87,14 @@ struct LayoutInfo {
LayoutInfo() = default;
LayoutInfo(const xegpu::DistributeLayoutAttr &layout) : storage(layout) {}
- // Two lattice values are equal if they have `some` layout. The actual
- // content of the layout does not matter.
+ // Two lattice values are equal if they are both unassigned, or both assigned
+ // with the same layout.
bool operator==(const LayoutInfo &other) const {
- return this->isAssigned() == other.isAssigned();
+ if (isAssigned() != other.isAssigned())
+ return false;
+ if (!isAssigned())
+ return true;
+ return storage.isEqualTo(other.storage);
}
static LayoutInfo meet(const LayoutInfo &lhs, const LayoutInfo &rhs);
@@ -141,6 +146,10 @@ void LayoutInfo::print(raw_ostream &os) const {
LayoutInfo LayoutInfo::meet(const LayoutInfo &lhs, const LayoutInfo &rhs) {
if (!lhs.isAssigned())
return rhs;
+ if (!rhs.isAssigned())
+ return lhs;
+ if (!lhs.isSliceLayout() && rhs.isSliceLayout())
+ return rhs;
return lhs;
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 515c59db72819..6bf7fff810006 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -653,3 +653,32 @@ func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vect
return
}
}
+
+// -----
+// A value with two consumers: one back-propagates a plain layout (the store_nd
+// of %trunc), the other back-propagates a slice layout (the broadcast/transpose
+// chain feeding the reduction result). `meet` must prefer the slice layout, and
+// the inst_data / lane_layout / lane_data fields must be preserved on it.
+gpu.module @test {
+ // CHECK-LABEL: truncf_prefers_slice
+ // CHECK: %[[TRUNC:.*]] = arith.truncf
+ // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [4, 8, 4], lane_layout = [4, 1, 4], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
+ // CHECK-SAME: : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
+ gpu.func @truncf_prefers_slice(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
+ %dst_plain: memref<32x4xf8E8M0FNU>) kernel {
+ %cst = arith.constant dense<0xFF80> : vector<32x4xbf16>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<32x128xbf16> -> !xegpu.tensor_desc<32x128xbf16>
+ %load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<32x128xbf16> -> vector<32x128xbf16>
+ %sc1 = vector.shape_cast %load : vector<32x128xbf16> to vector<32x4x32xbf16>
+ %red = vector.multi_reduction <maximumf>, %sc1, %cst [2] : vector<32x4x32xbf16> to vector<32x4xbf16>
+ %trunc = arith.truncf %red : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
+ %bcast = vector.broadcast %trunc : vector<32x4xf8E8M0FNU> to vector<32x32x4xf8E8M0FNU>
+ %bcast2 = vector.transpose %bcast, [1, 2, 0] : vector<32x32x4xf8E8M0FNU> to vector<32x4x32xf8E8M0FNU>
+ %sc2 = vector.shape_cast %bcast2 : vector<32x4x32xf8E8M0FNU> to vector<32x128xf8E8M0FNU>
+ %tdesc_red = xegpu.create_nd_tdesc %dst_red : memref<32x128xf8E8M0FNU> -> !xegpu.tensor_desc<32x128xf8E8M0FNU>
+ xegpu.store_nd %sc2, %tdesc_red[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<32x128xf8E8M0FNU>, !xegpu.tensor_desc<32x128xf8E8M0FNU>
+ %tdesc_plain = xegpu.create_nd_tdesc %dst_plain : memref<32x4xf8E8M0FNU> -> !xegpu.tensor_desc<32x4xf8E8M0FNU>
+ xegpu.store_nd %trunc, %tdesc_plain[0, 0] <{layout = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>}> : vector<32x4xf8E8M0FNU>, !xegpu.tensor_desc<32x4xf8E8M0FNU>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
index f584ede003bd4..9bfb783299055 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
@@ -8,7 +8,7 @@
// XFAIL: *
// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
-#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
#b_packed = #xegpu.layout<sg_layout = [2, 2], sg_data = [512, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
#b = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
#c = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
>From c31ad5f6b4f3d97c50c91e62737b39cd07a23ecf Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 01:55:14 +0000
Subject: [PATCH 2/9] [MLIR][XeGPU] Prefer nearer consumer's layout in
LayoutInfo::meet
When a value is demanded by multiple consumers during backward layout
propagation, prefer the layout of the consumer that is nearer to the
producer in program order. This tends to preserve a consumer's layout as
far up the def chain as possible, reducing layout conversions. It is a
hint, not an optimum.
Each op is assigned a program-order index via a pre-order walk (matching
printed-IR order), so a use inside an scf.for body is nearer than a use
after the loop. LayoutInfo carries the demanding op's index in a new
`programOrder` field, stamped by the single-argument constructor from a
file-scoped `currentProgramOrder` that visitOperation sets, so the ~30
meet() call sites are unchanged. `programOrder` is not propagated (each
visited op stamps its own index) and is excluded from operator==.
This replaces the earlier slice-layout preference in meet(), which is
now redundant: distinct users always have distinct indices, so program
order decides every real conflict. Also removes the dead
LayoutInfo::transpose declaration.
Updates the multiple-uses propagate-layout tests to reflect that the
nearer consumer's layout now wins.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 75 +++++++++++++++++--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 20 ++---
2 files changed, 77 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ff9239b637a10..72b10e1d3265f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -30,6 +30,7 @@
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
@@ -38,6 +39,7 @@
#include "llvm/Support/Debug.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
+#include <limits>
namespace mlir {
namespace xegpu {
@@ -58,6 +60,14 @@ namespace {
// LayoutInfo
//===----------------------------------------------------------------------===//
+/// Program-order index of the op currently being visited by the backward
+/// analysis. `visitOperation` sets this before dispatching, and the
+/// single-argument LayoutInfo constructor stamps it onto every demand pushed to
+/// an operand, so the ~30 `operand->meet(LayoutInfo(...))` call sites need no
+/// change. A larger index means farther from the producer; the sentinel max
+/// marks demands with no associated op (e.g. exit state).
+static int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
+
/// Helper class for tracking the analysis state of an mlir value. For layout
/// propagation, the analysis state is simply the distribution layout of
/// each value. The distribution layout information is encapsulated using
@@ -75,20 +85,34 @@ namespace {
/// assigned with the same layout.
/// 3) The meet operator works as follows:
/// - If only one side is assigned, return that side.
-/// - If both sides are assigned, prefer the side carrying a slice layout.
-/// If both (or neither) are slice layouts, prefer the lhs (current
-/// state) so an already assigned unique layout is not changed.
+/// - If both sides are assigned, prefer the layout demanded by the user
+/// that is nearer to the producer in program order (smaller
+/// `programOrder`); on a tie keep lhs.
+///
+/// The `programOrder` field records the program-order index of the consumer op
+/// that demanded the layout (see `currentProgramOrder`). During this backward
+/// analysis a value can be demanded by several users; keeping the nearest one
+/// tends to preserve a consumer's layout as far up the def chain as possible,
+/// minimizing layout conversions. This is a hint, not an optimum.
+/// `programOrder` is never propagated up the chain - each visited op stamps its
+/// own index - so it is deliberately excluded from `operator==`.
struct LayoutInfo {
private:
xegpu::DistributeLayoutAttr storage = nullptr;
+ // Program-order index of the consumer op that demanded this layout. Smaller
+ // means nearer to the producer. Unassigned/unknown demands sort last.
+ int64_t programOrder = std::numeric_limits<int64_t>::max();
public:
LayoutInfo() = default;
- LayoutInfo(const xegpu::DistributeLayoutAttr &layout) : storage(layout) {}
+ LayoutInfo(const xegpu::DistributeLayoutAttr &layout);
+ LayoutInfo(const xegpu::DistributeLayoutAttr &layout, int64_t programOrder)
+ : storage(layout), programOrder(programOrder) {}
// Two lattice values are equal if they are both unassigned, or both assigned
- // with the same layout.
+ // with the same layout. `programOrder` is intentionally excluded: it is not
+ // propagated, so a pure order refinement must not be reported as a change.
bool operator==(const LayoutInfo &other) const {
if (isAssigned() != other.isAssigned())
return false;
@@ -105,8 +129,6 @@ struct LayoutInfo {
bool isAssigned() const { return storage != nullptr; }
- LayoutInfo transpose(ArrayRef<int64_t> permutation) const;
-
SmallVector<int> getLaneLayout() const;
SmallVector<int> getLaneData() const;
@@ -135,6 +157,11 @@ struct LayoutInfo {
void set(const xegpu::DistributeLayoutAttr &layout) { storage = layout; }
};
+// Stamp every demand pushed by the current op with that op's program-order
+// index so `meet` can prefer the nearest consumer.
+LayoutInfo::LayoutInfo(const xegpu::DistributeLayoutAttr &layout)
+ : storage(layout), programOrder(currentProgramOrder) {}
+
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -148,7 +175,10 @@ LayoutInfo LayoutInfo::meet(const LayoutInfo &lhs, const LayoutInfo &rhs) {
return rhs;
if (!rhs.isAssigned())
return lhs;
- if (!lhs.isSliceLayout() && rhs.isSliceLayout())
+ // Prefer the demand from the user nearer to the producer in program order.
+ // Distinct users always have distinct indices, so this decides every
+ // real conflict; on a tie (same op, or both unknown) keep lhs.
+ if (rhs.programOrder < lhs.programOrder)
return rhs;
return lhs;
}
@@ -185,6 +215,15 @@ class LayoutInfoPropagation
private:
xegpu::LayoutKind layoutKind;
unsigned indexBitWidth;
+
+ // Program-order index of every op, built lazily on first use via a pre-order
+ // walk of the top-level module/function (matching printed-IR order). Used to
+ // tell which consumer of a value is nearer to its producer.
+ DenseMap<Operation *, int64_t> programOrder;
+ // Returns the program-order index of `op`, populating `programOrder` from
+ // `op`'s top-level ancestor on first call.
+ int64_t getProgramOrder(Operation *op);
+
void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -300,9 +339,29 @@ class LayoutInfoPropagation
};
} // namespace
+int64_t LayoutInfoPropagation::getProgramOrder(Operation *op) {
+ auto it = programOrder.find(op);
+ if (it != programOrder.end())
+ return it->second;
+ // First time we see this op's tree: number every op under its top-level
+ // ancestor in pre-order (i.e. printed-IR order). Nested ops (e.g. inside an
+ // scf.for body) get an index between their parent and the parent's next
+ // sibling, so a use inside a loop is "nearer" than a use after it.
+ Operation *root = op;
+ while (root->getParentOp())
+ root = root->getParentOp();
+ int64_t counter = 0;
+ root->walk<WalkOrder::PreOrder>(
+ [&](Operation *o) { programOrder[o] = counter++; });
+ return programOrder.lookup(op);
+}
+
LogicalResult LayoutInfoPropagation::visitOperation(
Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
+ // Stamp demands pushed by this op with its program-order index so `meet` can
+ // prefer the nearest consumer.
+ currentProgramOrder = getProgramOrder(op);
TypeSwitch<Operation *>(op)
.Case(
[&](xegpu::DpasOp dpasOp) { visitDpasOp(dpasOp, operands, results); })
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 1a741ba21e128..49006e1ef9507 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -219,8 +219,8 @@ func.func @scatter_ops_custom_perm_layout(%src: memref<256xf16>) {
gpu.module @test {
// CHECK-LABEL: func.func @scatter_ops_preserve_load_perm_layout(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} dense<12> : vector<16xindex>
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]]
// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16xf16>
// CHECK: %[[ADD_RES:.*]] = arith.addf %[[LOAD_VEC]], %[[LOAD_VEC]] {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} : vector<16xf16>
@@ -341,10 +341,10 @@ func.func @binary_op_one_use(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.
gpu.module @test {
// CHECK-LABEL: func.func @binary_op_multiple_uses(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
-// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>,
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {
-// CHECK: %[[T2:.*]] = arith.addf %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf16>
+// CHECK: %[[T2:.*]] = arith.addf %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} : vector<16x16xf16>
// CHECK: %[[T3:.*]] = xegpu.dpas %{{.*}}, %[[T2]] {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]>} : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
// CHECK-NEXT: xegpu.store_nd %[[T3]], %[[ARG2]][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]>>
// CHECK-NEXT: xegpu.store_nd %[[T2]], %[[ARG3]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -429,18 +429,18 @@ func.func @if_single_use(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tens
gpu.module @test {
// CHECK-LABEL: func.func @if_multiple_uses(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
-// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>,
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: i1, %[[ARG3:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
// CHECK-SAME: %[[ARG4:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {
// CHECK: %[[T1:.*]] = scf.if %[[ARG2]] -> (vector<16x16xf16>) {
-// CHECK-NEXT: %[[T3:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>
+// CHECK-NEXT: %[[T3:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
// CHECK-NEXT: scf.yield %[[T3]] : vector<16x16xf16>
// CHECK-NEXT: } else {
-// CHECK-NEXT: %[[T4:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>
+// CHECK-NEXT: %[[T4:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
// CHECK-NEXT: scf.yield %[[T4]] : vector<16x16xf16>
-// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
func.func @if_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor_desc<16x16xf16>, %arg2: i1, %arg3: !xegpu.tensor_desc<8x16xf32>, %arg4: !xegpu.tensor_desc<16x16xf16>) {
%0 = xegpu.load_nd %arg0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
%1 = scf.if %arg2 -> (vector<16x16xf16>) {
>From 3c4ada8b1f50d0d9c48f3c8543aa0662f1f2164c Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 01:58:24 +0000
Subject: [PATCH 3/9] [MLIR][XeGPU] Reframe propagate-layout test as
nearer-user preference
Rename truncf_prefers_slice to truncf_prefers_nearer_user and update its
comment: on this branch %trunc keeps the broadcast chain's slice layout
because that consumer is nearer in program order, not because meet()
prefers slice layouts.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Dialect/XeGPU/propagate-layout-inst-data.mlir | 12 ++++++------
1 file changed, 6 insertions(+), 6 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 6bf7fff810006..47471b890d552 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -655,16 +655,16 @@ func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vect
}
// -----
-// A value with two consumers: one back-propagates a plain layout (the store_nd
-// of %trunc), the other back-propagates a slice layout (the broadcast/transpose
-// chain feeding the reduction result). `meet` must prefer the slice layout, and
-// the inst_data / lane_layout / lane_data fields must be preserved on it.
+// %trunc has two consumers: the broadcast/transpose chain (nearer in program
+// order, back-propagates a slice layout) and the store_nd of %trunc (farther,
+// back-propagates a plain layout). `meet` must keep the nearer consumer's
+// layout, preserving its inst_data / lane_layout / lane_data fields.
gpu.module @test {
- // CHECK-LABEL: truncf_prefers_slice
+ // CHECK-LABEL: truncf_prefers_nearer_user
// CHECK: %[[TRUNC:.*]] = arith.truncf
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [4, 8, 4], lane_layout = [4, 1, 4], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
// CHECK-SAME: : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
- gpu.func @truncf_prefers_slice(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
+ gpu.func @truncf_prefers_nearer_user(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
%dst_plain: memref<32x4xf8E8M0FNU>) kernel {
%cst = arith.constant dense<0xFF80> : vector<32x4xbf16>
%tdesc = xegpu.create_nd_tdesc %src : memref<32x128xbf16> -> !xegpu.tensor_desc<32x128xbf16>
>From 9c4de99b6094f502c82dcd411b8eca2c1c18141d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 02:00:09 +0000
Subject: [PATCH 4/9] git format
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 72b10e1d3265f..75e8db468f361 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -320,18 +320,18 @@ class LayoutInfoPropagation
visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) override;
- void visitBranchOperand(OpOperand &operand) override {};
+ void visitBranchOperand(OpOperand &operand) override{};
- void visitCallOperand(OpOperand &operand) override {};
+ void visitCallOperand(OpOperand &operand) override{};
void
visitNonControlFlowArguments(RegionSuccessor &successor,
- ArrayRef<BlockArgument> arguments) override {};
+ ArrayRef<BlockArgument> arguments) override{};
- void visitExternalCall(CallOpInterface call,
- ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override {
- };
+ void
+ visitExternalCall(CallOpInterface call,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) override{};
void setToExitState(LayoutInfoLattice *lattice) override {
(void)lattice->meet(LayoutInfo());
>From 1db9aa8f8f3ccd82194b591fc3f4754a056c416d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 02:11:28 +0000
Subject: [PATCH 5/9] [MLIR][XeGPU] clang-format XeGPUPropagateLayout.cpp
Apply clang-format fixes flagged by CI (override{} -> override {}) on the
empty visitor overrides touched by this change.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 75e8db468f361..72b10e1d3265f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -320,18 +320,18 @@ class LayoutInfoPropagation
visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) override;
- void visitBranchOperand(OpOperand &operand) override{};
+ void visitBranchOperand(OpOperand &operand) override {};
- void visitCallOperand(OpOperand &operand) override{};
+ void visitCallOperand(OpOperand &operand) override {};
void
visitNonControlFlowArguments(RegionSuccessor &successor,
- ArrayRef<BlockArgument> arguments) override{};
+ ArrayRef<BlockArgument> arguments) override {};
- void
- visitExternalCall(CallOpInterface call,
- ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override{};
+ void visitExternalCall(CallOpInterface call,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) override {
+ };
void setToExitState(LayoutInfoLattice *lattice) override {
(void)lattice->meet(LayoutInfo());
>From a6ee4af226d56c74c108d08ab300d5aa318308e9 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 10 Jul 2026 19:24:55 +0000
Subject: [PATCH 6/9] [MLIR][XeGPU] Address review: replace global
currentProgramOrder with member helper
Move currentProgramOrder from a file-scoped global into LayoutInfoPropagation
and add a private makeLayoutInfo helper that stamps it, so the LayoutInfo
constructor no longer depends on implicit global mutable state. Rename the
single-arg LayoutInfo(...) call sites to makeLayoutInfo(...) and drop the now
unused single-arg constructor. Also fix the stale operator== comment to
describe each branch.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 194 ++++++++----------
1 file changed, 82 insertions(+), 112 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index a5618b5f84ed3..c73749128906b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -60,14 +60,6 @@ namespace {
// LayoutInfo
//===----------------------------------------------------------------------===//
-/// Program-order index of the op currently being visited by the backward
-/// analysis. `visitOperation` sets this before dispatching, and the
-/// single-argument LayoutInfo constructor stamps it onto every demand pushed to
-/// an operand, so the ~30 `operand->meet(LayoutInfo(...))` call sites need no
-/// change. A larger index means farther from the producer; the sentinel max
-/// marks demands with no associated op (e.g. exit state).
-static int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
-
/// Helper class for tracking the analysis state of an mlir value. For layout
/// propagation, the analysis state is simply the distribution layout of
/// each value. The distribution layout information is encapsulated using
@@ -90,12 +82,14 @@ static int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
/// `programOrder`); on a tie keep lhs.
///
/// The `programOrder` field records the program-order index of the consumer op
-/// that demanded the layout (see `currentProgramOrder`). During this backward
-/// analysis a value can be demanded by several users; keeping the nearest one
-/// tends to preserve a consumer's layout as far up the def chain as possible,
-/// minimizing layout conversions. This is a hint, not an optimum.
-/// `programOrder` is never propagated up the chain - each visited op stamps its
-/// own index - so it is deliberately excluded from `operator==`.
+/// that demanded the layout (stamped via
+/// `LayoutInfoPropagation::makeLayoutInfo` from
+/// `LayoutInfoPropagation::currentProgramOrder`). During this backward analysis
+/// a value can be demanded by several users; keeping the nearest one tends to
+/// preserve a consumer's layout as far up the def chain as possible, minimizing
+/// layout conversions. This is a hint, not an optimum. `programOrder` is never
+/// propagated up the chain - each visited op stamps its own index - so it is
+/// deliberately excluded from `operator==`.
struct LayoutInfo {
private:
@@ -106,13 +100,15 @@ struct LayoutInfo {
public:
LayoutInfo() = default;
- LayoutInfo(const xegpu::DistributeLayoutAttr &layout);
LayoutInfo(const xegpu::DistributeLayoutAttr &layout, int64_t programOrder)
: storage(layout), programOrder(programOrder) {}
- // Two lattice values are equal if they are both unassigned, or both assigned
- // with the same layout. `programOrder` is intentionally excluded: it is not
- // propagated, so a pure order refinement must not be reported as a change.
+ // Equality by assignment state and, when both assigned, by the layout:
+ // - one assigned, the other not -> not equal;
+ // - both unassigned -> equal;
+ // - both assigned -> equal iff the layouts match.
+ // `programOrder` is intentionally excluded: it is not propagated, so a pure
+ // order refinement must not be reported as a change.
bool operator==(const LayoutInfo &other) const {
if (isAssigned() != other.isAssigned())
return false;
@@ -157,11 +153,6 @@ struct LayoutInfo {
void set(const xegpu::DistributeLayoutAttr &layout) { storage = layout; }
};
-// Stamp every demand pushed by the current op with that op's program-order
-// index so `meet` can prefer the nearest consumer.
-LayoutInfo::LayoutInfo(const xegpu::DistributeLayoutAttr &layout)
- : storage(layout), programOrder(currentProgramOrder) {}
-
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -224,6 +215,11 @@ class LayoutInfoPropagation
// `op`'s top-level ancestor on first call.
int64_t getProgramOrder(Operation *op);
+ int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
+ LayoutInfo makeLayoutInfo(const xegpu::DistributeLayoutAttr &layout) {
+ return LayoutInfo(layout, currentProgramOrder);
+ }
+
void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -459,42 +455,6 @@ bool LayoutInfoPropagation::hasParamsOfLayoutKind(
return false;
}
-// This function returns all layouts for the given sgCount, whose sgData:
-// 1. Evenly divides the wgShape.
-// 2. Is a multiple of instData.
-// Example:
-// wgShape = [128, 64], instData = [8, 16], sgCount = 32
-// Returns layouts:
-// [(8,4), (16,2)], which correspond to sgData [16,16] and [8,32].
-SmallVector<std::pair<int, int>>
-getSgLayoutCandidates(ArrayRef<int64_t> wgShape, ArrayRef<int> instData,
- int64_t sgCount) {
- SmallVector<std::pair<int, int>> candidates;
- for (int sgLayout0 = 1; sgLayout0 <= sgCount; ++sgLayout0) {
- if (sgCount % sgLayout0)
- continue;
- int sgLayout1 = sgCount / sgLayout0;
- int sgData0 = wgShape[0] / sgLayout0;
- int sgData1 = wgShape[1] / sgLayout1;
- if ((wgShape[0] % sgLayout0 || wgShape[1] % sgLayout1) ||
- (sgData0 % instData[0] || sgData1 % instData[1]))
- continue;
- candidates.emplace_back(sgLayout0, sgLayout1);
- }
- // Sort primarily by how balanced they are
- // (i.e., minimize the absolute difference between the two dimensions), and
- // secondarily by the first dimension in ascending order.
- llvm::sort(candidates, [](const std::pair<int, int> &lhs,
- const std::pair<int, int> &rhs) {
- int diffLhs = std::abs(lhs.first - lhs.second);
- int diffRhs = std::abs(rhs.first - rhs.second);
- if (diffLhs != diffRhs)
- return diffLhs < diffRhs;
- return lhs.first < rhs.first;
- });
- return candidates;
-}
-
FailureOr<int64_t>
getNumSg(Operation *op, const int sgSize,
xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
@@ -525,7 +485,7 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
return;
xegpu::DistributeLayoutAttr anchorLayout = prefetch.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
- prefetchLayout = LayoutInfo(anchorLayout);
+ prefetchLayout = makeLayoutInfo(anchorLayout);
if (layoutKind == xegpu::LayoutKind::InstData) {
const auto *uArchInstruction =
dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(
@@ -542,7 +502,7 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
return;
}
prefetch.setLayoutAttr(*completed);
- prefetchLayout = LayoutInfo(*completed);
+ prefetchLayout = makeLayoutInfo(*completed);
}
} else {
auto tdescTy = prefetch.getTensorDescType();
@@ -560,7 +520,7 @@ void LayoutInfoPropagation::visitPrefetchNdOp(
"Failed to determine required layout for prefetch_nd.");
return;
}
- prefetchLayout = LayoutInfo(layoutAttr);
+ prefetchLayout = makeLayoutInfo(layoutAttr);
prefetch.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the source tensor descriptor.
@@ -614,10 +574,11 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
auto srcLayoutAttr = xegpu::inferMultiReductionSourceLayout(
requiredResLayoutAttr, reductionDims);
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
// Accumulator should have the same layout as the result.
propagateIfChanged(operands[1],
- operands[1]->meet(LayoutInfo(requiredResLayoutAttr)));
+ operands[1]->meet(makeLayoutInfo(requiredResLayoutAttr)));
}
void LayoutInfoPropagation::visitVectorReductionOp(
@@ -635,10 +596,11 @@ void LayoutInfoPropagation::visitVectorReductionOp(
xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
auto srcLayoutAttr = xegpu::inferReductionSourceLayout(requiredResLayoutAttr);
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
if (reduction.getAcc())
- propagateIfChanged(operands[1],
- operands[1]->meet(LayoutInfo(requiredResLayoutAttr)));
+ propagateIfChanged(
+ operands[1], operands[1]->meet(makeLayoutInfo(requiredResLayoutAttr)));
}
void LayoutInfoPropagation::visitVectorBroadCastOp(
@@ -665,7 +627,8 @@ void LayoutInfoPropagation::visitVectorBroadCastOp(
xegpu::DistributeLayoutAttr srcLayoutAttr =
xegpu::inferBroadcastSourceLayout(resultLayoutAttr, resShape, srcShape);
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
}
void LayoutInfoPropagation::visitShapeCastOp(
@@ -690,7 +653,8 @@ void LayoutInfoPropagation::visitShapeCastOp(
return;
}
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
}
/// Set the layouts for DPAS A, B, and C operands.
@@ -716,9 +680,9 @@ void LayoutInfoPropagation::visitDpasOp(
"Expected anchor layout for DPAS A operand.");
assert(hasParamsOfLayoutKind(anchorLayoutB) &&
"Expected anchor layout for DPAS B operand.");
- dpasALayout = LayoutInfo(anchorLayoutA);
- dpasBLayout = LayoutInfo(anchorLayoutB);
- dpasCDLayout = LayoutInfo(anchorLayoutCD);
+ dpasALayout = makeLayoutInfo(anchorLayoutA);
+ dpasBLayout = makeLayoutInfo(anchorLayoutB);
+ dpasCDLayout = makeLayoutInfo(anchorLayoutCD);
if (layoutKind == xegpu::LayoutKind::InstData) {
auto completed = xegpu::completeDpasLaneLayoutFromInstData(
anchorLayoutA, anchorLayoutB, anchorLayoutCD, aTy, bTy, cdTy, uArch);
@@ -731,9 +695,9 @@ void LayoutInfoPropagation::visitDpasOp(
dpas.setLayoutAAttr(completedA);
dpas.setLayoutBAttr(completedB);
dpas.setLayoutCdAttr(completedCD);
- dpasALayout = LayoutInfo(completedA);
- dpasBLayout = LayoutInfo(completedB);
- dpasCDLayout = LayoutInfo(completedCD);
+ dpasALayout = makeLayoutInfo(completedA);
+ dpasBLayout = makeLayoutInfo(completedB);
+ dpasCDLayout = makeLayoutInfo(completedCD);
}
} else {
@@ -769,9 +733,9 @@ void LayoutInfoPropagation::visitDpasOp(
dpas.setLayoutAAttr(requiredALayout);
dpas.setLayoutBAttr(requiredBLayout);
dpas.setLayoutCdAttr(requiredCDLayoutAttr);
- dpasALayout = LayoutInfo(requiredALayout);
- dpasBLayout = LayoutInfo(requiredBLayout);
- dpasCDLayout = LayoutInfo(requiredCDLayoutAttr);
+ dpasALayout = makeLayoutInfo(requiredALayout);
+ dpasBLayout = makeLayoutInfo(requiredBLayout);
+ dpasCDLayout = makeLayoutInfo(requiredCDLayoutAttr);
}
propagateIfChanged(operands[0], operands[0]->meet(dpasALayout));
propagateIfChanged(operands[1], operands[1]->meet(dpasBLayout));
@@ -818,9 +782,9 @@ void LayoutInfoPropagation::visitDpasMxOp(
hasParamsOfLayoutKind(anchorLayoutA) &&
hasParamsOfLayoutKind(anchorLayoutB) &&
hasParamsOfLayoutKind(anchorLayoutCD)) {
- dpasMxALayout = LayoutInfo(anchorLayoutA);
- dpasMxBLayout = LayoutInfo(anchorLayoutB);
- dpasMxCDLayout = LayoutInfo(anchorLayoutCD);
+ dpasMxALayout = makeLayoutInfo(anchorLayoutA);
+ dpasMxBLayout = makeLayoutInfo(anchorLayoutB);
+ dpasMxCDLayout = makeLayoutInfo(anchorLayoutCD);
// Get scale layouts if available
xegpu::DistributeLayoutAttr anchorLayoutAScale =
@@ -828,9 +792,9 @@ void LayoutInfoPropagation::visitDpasMxOp(
xegpu::DistributeLayoutAttr anchorLayoutBScale =
dpasMx.getLayoutBScaleAttr();
if (anchorLayoutAScale)
- dpasMxAScaleLayout = LayoutInfo(anchorLayoutAScale);
+ dpasMxAScaleLayout = makeLayoutInfo(anchorLayoutAScale);
if (anchorLayoutBScale)
- dpasMxBScaleLayout = LayoutInfo(anchorLayoutBScale);
+ dpasMxBScaleLayout = makeLayoutInfo(anchorLayoutBScale);
if (layoutKind == xegpu::LayoutKind::InstData) {
auto completed = xegpu::completeDpasMxLaneLayoutFromInstData(
@@ -846,16 +810,16 @@ void LayoutInfoPropagation::visitDpasMxOp(
dpasMx.setLayoutAAttr(completedA);
dpasMx.setLayoutBAttr(completedB);
dpasMx.setLayoutCdAttr(completedCD);
- dpasMxALayout = LayoutInfo(completedA);
- dpasMxBLayout = LayoutInfo(completedB);
- dpasMxCDLayout = LayoutInfo(completedCD);
+ dpasMxALayout = makeLayoutInfo(completedA);
+ dpasMxBLayout = makeLayoutInfo(completedB);
+ dpasMxCDLayout = makeLayoutInfo(completedCD);
if (completedAScale) {
dpasMx.setLayoutAScaleAttr(completedAScale);
- dpasMxAScaleLayout = LayoutInfo(completedAScale);
+ dpasMxAScaleLayout = makeLayoutInfo(completedAScale);
}
if (completedBScale) {
dpasMx.setLayoutBScaleAttr(completedBScale);
- dpasMxBScaleLayout = LayoutInfo(completedBScale);
+ dpasMxBScaleLayout = makeLayoutInfo(completedBScale);
}
}
} else {
@@ -897,13 +861,13 @@ void LayoutInfoPropagation::visitDpasMxOp(
if (requiredBScaleLayout)
dpasMx.setLayoutBScaleAttr(requiredBScaleLayout);
- dpasMxALayout = LayoutInfo(requiredALayout);
- dpasMxBLayout = LayoutInfo(requiredBLayout);
- dpasMxCDLayout = LayoutInfo(requiredCDLayoutAttr);
+ dpasMxALayout = makeLayoutInfo(requiredALayout);
+ dpasMxBLayout = makeLayoutInfo(requiredBLayout);
+ dpasMxCDLayout = makeLayoutInfo(requiredCDLayoutAttr);
if (requiredAScaleLayout)
- dpasMxAScaleLayout = LayoutInfo(requiredAScaleLayout);
+ dpasMxAScaleLayout = makeLayoutInfo(requiredAScaleLayout);
if (requiredBScaleLayout)
- dpasMxBScaleLayout = LayoutInfo(requiredBScaleLayout);
+ dpasMxBScaleLayout = makeLayoutInfo(requiredBScaleLayout);
}
// Propagate layouts to operands. Because acc, scale_a, scale_b are all
@@ -941,7 +905,7 @@ void LayoutInfoPropagation::visitStoreNdOp(
return;
xegpu::DistributeLayoutAttr anchorLayout = store.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
- storeLayout = LayoutInfo(anchorLayout);
+ storeLayout = makeLayoutInfo(anchorLayout);
if (layoutKind == xegpu::LayoutKind::InstData) {
const auto *uArchInstruction =
@@ -959,7 +923,7 @@ void LayoutInfoPropagation::visitStoreNdOp(
return;
}
store.setLayoutAttr(*completed);
- storeLayout = LayoutInfo(*completed);
+ storeLayout = makeLayoutInfo(*completed);
}
} else {
auto numSgOrErr = getNumSg(store, uArch->getSubgroupSize());
@@ -975,7 +939,7 @@ void LayoutInfoPropagation::visitStoreNdOp(
store.emitWarning("Failed to determine required layout for store_nd.");
return;
}
- storeLayout = LayoutInfo(layoutAttr);
+ storeLayout = makeLayoutInfo(layoutAttr);
store.setLayoutAttr(layoutAttr);
}
// Propagate the layout to the value operand.
@@ -1001,7 +965,7 @@ void LayoutInfoPropagation::visitLoadNdOp(
dyn_cast<xegpu::DistributeLayoutAttr>(valueLayout.get());
xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();
if (hasParamsOfLayoutKind(anchorLayout)) {
- loadLayout = LayoutInfo(anchorLayout);
+ loadLayout = makeLayoutInfo(anchorLayout);
if (layoutKind == xegpu::LayoutKind::InstData &&
!consumerLayoutAttr.getEffectiveLaneLayoutAsInt().empty()) {
const auto *uArchInstruction =
@@ -1019,7 +983,7 @@ void LayoutInfoPropagation::visitLoadNdOp(
return;
}
load.setLayoutAttr(*completed);
- loadLayout = LayoutInfo(*completed);
+ loadLayout = makeLayoutInfo(*completed);
}
} else {
auto numSgOrErr =
@@ -1036,7 +1000,7 @@ void LayoutInfoPropagation::visitLoadNdOp(
load.emitWarning("Failed to determine required layout for load_nd.");
return;
}
- loadLayout = LayoutInfo(layoutAttr);
+ loadLayout = makeLayoutInfo(layoutAttr);
load.setLayoutAttr(layoutAttr);
}
// Propagate the new layout to the tensor descriptor operand.
@@ -1092,7 +1056,7 @@ void LayoutInfoPropagation::visitConvertLayoutOp(
}
xegpu::DistributeLayoutAttr anchorLayout = convert.getInputLayoutAttr();
- LayoutInfo convertLayout(anchorLayout);
+ LayoutInfo convertLayout = makeLayoutInfo(anchorLayout);
// Propagate the new layout to the tensor descriptor operand.
propagateIfChanged(operands[0], operands[0]->meet(convertLayout));
}
@@ -1113,7 +1077,8 @@ void LayoutInfoPropagation::visitTransposeOp(
consumerLayoutAttr, transpose.getPermutation());
// Propagate the new layout to the vector operand.
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
}
/// For vector::BitCastOp, the lane_data of the source layout is changed based
@@ -1147,7 +1112,8 @@ void LayoutInfoPropagation::visitVectorBitcastOp(
auto srcLayoutAttr = xegpu::inferBitCastSourceLayout(
requiredResLayoutAttr, outElemTyBitWidth, inElemTyBitWidth);
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
}
/// For vector::InterleaveOp, the result has double the innermost dimension
@@ -1182,8 +1148,10 @@ void LayoutInfoPropagation::visitVectorInterleaveOp(
xegpu::inferInterleaveSourceLayout(requiredResLayoutAttr);
// Both operands (lhs and rhs) get the same source layout
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
- propagateIfChanged(operands[1], operands[1]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[1],
+ operands[1]->meet(makeLayoutInfo(srcLayoutAttr)));
}
/// For vector::DeinterleaveOp, the source has double the innermost dimension
@@ -1205,7 +1173,8 @@ void LayoutInfoPropagation::visitVectorDeinterleaveOp(
// dim) No setup function needed - just infer directly
auto srcLayoutAttr = xegpu::inferDeinterleaveSourceLayout(consumerLayoutAttr);
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
}
void LayoutInfoPropagation::visitInsertStridedSliceOp(
@@ -1234,9 +1203,10 @@ void LayoutInfoPropagation::visitInsertStridedSliceOp(
auto srcLayoutAttr = xegpu::inferInsertStridedSliceSourceLayout(
requiredResLayoutAttr, resVecType.getShape(), srcVecType.getShape());
- propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
+ propagateIfChanged(operands[0],
+ operands[0]->meet(makeLayoutInfo(srcLayoutAttr)));
propagateIfChanged(operands[1],
- operands[1]->meet(LayoutInfo(requiredResLayoutAttr)));
+ operands[1]->meet(makeLayoutInfo(requiredResLayoutAttr)));
}
/// Propagate the layout of the result to the tensor descriptor, mask and
@@ -1291,8 +1261,8 @@ void LayoutInfoPropagation::visitLoadGatherOp(
assert((chunkSize <= 1) || (layoutKind != xegpu::LayoutKind::Subgroup));
auto maskLayoutAttr = xegpu::inferMaskOffsetLayoutForScatterIO(
requiredAnchorLayoutAttr, chunkSize);
- LayoutInfo maskLayoutInfo = LayoutInfo(maskLayoutAttr);
- auto loadLayoutInfo = LayoutInfo(requiredAnchorLayoutAttr);
+ LayoutInfo maskLayoutInfo = makeLayoutInfo(maskLayoutAttr);
+ auto loadLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
// Propagate the new layout to the tensor descriptor operand.
if (isa<xegpu::TensorDescType>(load.getSourceType()))
@@ -1357,11 +1327,11 @@ void LayoutInfoPropagation::visitStoreScatterOp(
storeScatter.setLayoutAttr(requiredAnchorLayoutAttr);
}
- LayoutInfo srcLayoutInfo = LayoutInfo(requiredAnchorLayoutAttr);
+ LayoutInfo srcLayoutInfo = makeLayoutInfo(requiredAnchorLayoutAttr);
assert((chunkSize <= 1) || (layoutKind != xegpu::LayoutKind::Subgroup));
auto maskLayoutAttr = xegpu::inferMaskOffsetLayoutForScatterIO(
requiredAnchorLayoutAttr, chunkSize);
- LayoutInfo maskLayoutInfo = LayoutInfo(maskLayoutAttr);
+ LayoutInfo maskLayoutInfo = makeLayoutInfo(maskLayoutAttr);
// Propagate the payload operand layout
propagateIfChanged(operands[0], operands[0]->meet(srcLayoutInfo));
@@ -1451,7 +1421,7 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
}
storeMatrix.setLayoutAttr(requiredAnchorLayoutAttr);
}
- layout = LayoutInfo(requiredAnchorLayoutAttr);
+ layout = makeLayoutInfo(requiredAnchorLayoutAttr);
propagateIfChanged(operands[0], operands[0]->meet(layout));
}
>From e19b889e3d33f34618d7d2b65eb8939956b0f080 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 10 Jul 2026 19:26:12 +0000
Subject: [PATCH 7/9] [MLIR][XeGPU] Update test expectations for nearer-user
layout preference
Update the truncf_prefers_nearer_user CHECK line and the simple_mxfp_gemm
quantizeA layout to match the layouts produced by preferring the nearer
consumer's demand in LayoutInfo::meet.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir | 2 +-
.../Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 1f0acc68b131e..e8cddb8fa1de3 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -653,7 +653,7 @@ func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vect
gpu.module @test {
// CHECK-LABEL: truncf_prefers_nearer_user
// CHECK: %[[TRUNC:.*]] = arith.truncf
- // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [4, 8, 4], lane_layout = [4, 1, 4], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
+ // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 8, 1], lane_layout = [16, 1, 1], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
// CHECK-SAME: : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
gpu.func @truncf_prefers_nearer_user(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
%dst_plain: memref<32x4xf8E8M0FNU>) kernel {
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
index cfaa8077c742c..87603594db5bf 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
@@ -8,7 +8,7 @@
// XFAIL: *
// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
-#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
#a_ld = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
#b_packed = #xegpu.layout<sg_layout = [2, 2], sg_data = [512, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
#b = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
>From 557c0167940ef4dacfc64c010e4122e5a20cb7c9 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 10 Jul 2026 23:11:09 +0000
Subject: [PATCH 8/9] [MLIR][XeGPU] Fix overstated programOrder comments
The comments claimed excluding programOrder from operator== was required
for convergence. It is not: order decreases monotonically under meet, so
including it would only cause redundant re-propagation, not incorrect
results. Drop the misleading wording.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 4 +---
1 file changed, 1 insertion(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index c73749128906b..1681d295ae0ff 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -89,7 +89,7 @@ namespace {
/// preserve a consumer's layout as far up the def chain as possible, minimizing
/// layout conversions. This is a hint, not an optimum. `programOrder` is never
/// propagated up the chain - each visited op stamps its own index - so it is
-/// deliberately excluded from `operator==`.
+/// excluded from `operator==`.
struct LayoutInfo {
private:
@@ -107,8 +107,6 @@ struct LayoutInfo {
// - one assigned, the other not -> not equal;
// - both unassigned -> equal;
// - both assigned -> equal iff the layouts match.
- // `programOrder` is intentionally excluded: it is not propagated, so a pure
- // order refinement must not be reported as a change.
bool operator==(const LayoutInfo &other) const {
if (isAssigned() != other.isAssigned())
return false;
>From f45cf2a99627e236189da51344a6f763cbfbb46b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 11 Jul 2026 15:26:24 +0000
Subject: [PATCH 9/9] [MLIR][XeGPU] Add SLM privatization to WG-to-SG
distribution
Demote a shared local memory (space 3) buffer to subgroup-private memory
(space 4) when every load_matrix/store_matrix on it uses the same sg_layout,
sg_data, data shape and offsets, and the workgroup tile is evenly distributed
across subgroups. In that case each subgroup accesses an identical,
non-overlapping region, so the buffer and mem_desc are shrunk to the
per-subgroup size and matrix ops are re-indexed with local, subgroup-id-free
offsets.
Co-Authored-By: Claude Opus 4.8 <noreply at anthropic.com>
---
.../mlir/Dialect/XeGPU/IR/XeGPUDialect.td | 5 +
.../mlir/Dialect/XeGPU/IR/XeGPUTypes.td | 11 +-
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 11 +
.../Transforms/XeGPUWgToSgDistribute.cpp | 308 +++++++++++++++++-
mlir/test/Dialect/XeGPU/invalid.mlir | 2 +-
.../XeGPU/xegpu-wg-to-sg-privatize.mlir | 207 ++++++++++++
6 files changed, 529 insertions(+), 15 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-privatize.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
index b1490c7742a26..34d38844051d1 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
@@ -43,6 +43,11 @@ def XeGPU_Dialect : Dialect {
/// xevm::AddrSpace::SHARED, or a GPU workgroup memory address space.
static bool isSharedMemory(const MemRefType &memrefTy);
+ /// Checks if the given memref type represents subgroup-private memory.
+ /// Returns true if the memory space is address space 4, or a GPU private
+ /// memory address space.
+ static bool isPrivateMemory(const MemRefType &memrefTy);
+
/// drops/slices the shape in the specified dims, and return the rest. e.g.,
/// for shape = [32, 64, 8], dims = [0, 2], it will return [64]
template<typename T, typename U>
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
index 0423303c23493..0c2acc550d93a 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
@@ -40,14 +40,17 @@ class XeGPUTypeDef<string name, string typeMnemonic, list<Trait> traits = [],
}
def isSharedPred : CPred<"XeGPUDialect::isSharedMemory(llvm::cast<mlir::MemRefType>($_self))">;
+def isSharedOrPrivatePred : CPred<[{
+ XeGPUDialect::isSharedMemory(llvm::cast<mlir::MemRefType>($_self)) ||
+ XeGPUDialect::isPrivateMemory(llvm::cast<mlir::MemRefType>($_self))}]>;
class StaticShared1DMemRefOf<list<Type> allowedTypes> :
- ConfinedType<MemRefRankOf<allowedTypes, [1]>, [HasStaticShapePred, isSharedPred],
- "reside in share memory and statically 1d shaped " # MemRefOf<allowedTypes>.summary # " ",
+ ConfinedType<MemRefRankOf<allowedTypes, [1]>, [HasStaticShapePred, isSharedOrPrivatePred],
+ "reside in shared or private memory and statically 1d shaped " # MemRefOf<allowedTypes>.summary # " ",
"mlir::MemRefType">;
class StaticShared2DMemRefOf<list<Type> allowedTypes>:
- ConfinedType<MemRefRankOf<allowedTypes, [2]>, [HasStaticShapePred, isSharedPred],
- "reside in share memory and statically 2d shaped " # MemRefOf<allowedTypes>.summary # " ",
+ ConfinedType<MemRefRankOf<allowedTypes, [2]>, [HasStaticShapePred, isSharedOrPrivatePred],
+ "reside in shared or private memory and statically 2d shaped " # MemRefOf<allowedTypes>.summary # " ",
"mlir::MemRefType">;
def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 0951350e2fe82..1b1deb1edcf4f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -194,6 +194,17 @@ bool XeGPUDialect::isSharedMemory(const MemRefType &memrefTy) {
return gpu::GPUDialect::isWorkgroupMemoryAddressSpace(attr);
}
+bool XeGPUDialect::isPrivateMemory(const MemRefType &memrefTy) {
+ Attribute attr = memrefTy.getMemorySpace();
+ if (!attr)
+ return false; // Default memory space is not private memory
+ if (auto intAttr = llvm::dyn_cast_if_present<IntegerAttr>(attr))
+ return intAttr.getInt() == 4;
+ if (auto gpuAttr = llvm::dyn_cast_if_present<gpu::AddressSpaceAttr>(attr))
+ return gpuAttr.getValue() == gpu::GPUDialect::getPrivateAddressSpace();
+ return false;
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_BlockTensorDescAttr
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 9ca6b3c2b0272..bc540708264fc 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -36,6 +36,12 @@ using namespace mlir;
namespace {
+// Marker attribute carrying the common workgroup layout onto create_mem_desc so
+// the conversion pattern can shrink the private buffer to the per-subgroup size.
+// Set by the SLM privatization pre-phase and consumed by WgToSgCreateMemDescOp.
+static constexpr StringLiteral kPrivatizeLayoutAttrName =
+ "__xegpu_privatize_layout__";
+
// Retrieve the RangeAttr if it is specified.
static xegpu::RangeAttr getRangeSpecAttr(Operation *op) {
Operation *parent = op->getParentOfType<scf::IfOp>();
@@ -887,28 +893,76 @@ struct WgToSgStoreScatterOp
}
};
+// Returns true if `memDesc` was produced by a create_mem_desc whose source
+// memref lives in private memory (i.e. it was privatized by the pre-phase).
+static bool isPrivateMemDesc(Value memDesc) {
+ auto createOp = memDesc.getDefiningOp<xegpu::CreateMemDescOp>();
+ if (!createOp)
+ return false;
+ auto srcTy = dyn_cast<MemRefType>(createOp.getSource().getType());
+ return srcTy && xegpu::XeGPUDialect::isPrivateMemory(srcTy);
+}
+
+// Computes the per-subgroup local offset lists into a privatized (shrunk)
+// mem_desc. Because every subgroup owns a private copy sized to the distributed
+// shape, offsets are subgroup-id independent: round `r` maps to `r * sg_data`.
+static SmallVector<SmallVector<OpFoldResult>>
+genPrivateOffsetsList(ConversionPatternRewriter &rewriter, Location loc,
+ ArrayRef<int64_t> wgShape,
+ xegpu::DistributeLayoutAttr layout) {
+ SmallVector<int64_t> sgData = layout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> sgLayout = layout.getEffectiveSgLayoutAsInt();
+
+ // The private buffer holds the per-subgroup distributed shape, which spans all
+ // distribution rounds stacked along each dim (roundShape[d] = wgShape[d] /
+ // sg_layout[d] = sg_data[d] * rounds). This matches the shape allocated by
+ // WgToSgCreateMemDescOp, so stepping by sg_data yields in-bounds local offsets.
+ SmallVector<int64_t> roundShape(wgShape.size());
+ for (auto [i, dim] : llvm::enumerate(wgShape))
+ roundShape[i] = dim / sgLayout[i];
+
+ SmallVector<SmallVector<OpFoldResult>> offsetsList;
+ for (SmallVector<int64_t> off : StaticTileOffsetRange(roundShape, sgData)) {
+ SmallVector<OpFoldResult> offsets;
+ for (int64_t o : off)
+ offsets.push_back(rewriter.getIndexAttr(o));
+ offsetsList.push_back(std::move(offsets));
+ }
+ return offsetsList;
+}
+
struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {
using OpConversionPattern<xegpu::LoadMatrixOp>::OpConversionPattern;
LogicalResult
matchAndRewrite(xegpu::LoadMatrixOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- SmallVector<SmallVector<OpFoldResult>> offsetsList;
- if (failed(genOffsetsList(rewriter, op, offsetsList)))
- return failure();
-
ArrayRef<int64_t> wgShape = op.getDataShape();
VectorType valueTy = llvm::dyn_cast<VectorType>(op.getRes().getType());
assert(valueTy && "the value type must be vector type!");
Type elemTy = valueTy.getElementType();
xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
+
+ // Privatized buffers are shrunk to the per-subgroup size and indexed with
+ // local, subgroup-id-free offsets; SLM buffers keep the sg-relative offsets.
+ bool isPrivate = isPrivateMemDesc(op.getMemDesc());
+ auto memDesc = cast<TypedValue<xegpu::MemDescType>>(
+ isPrivate ? adaptor.getMemDesc()[0] : op.getMemDesc());
+ SmallVector<SmallVector<OpFoldResult>> offsetsList;
+ if (isPrivate) {
+ offsetsList = genPrivateOffsetsList(rewriter, op.getLoc(), wgShape,
+ layout);
+ } else if (failed(genOffsetsList(rewriter, op, offsetsList))) {
+ return failure();
+ }
+
SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;
VectorType newResTy = VectorType::get(sgShape, elemTy);
SmallVector<Value> newOps;
for (auto offsets : offsetsList) {
auto newOp = xegpu::LoadMatrixOp::create(rewriter, op.getLoc(), newResTy,
- op.getMemDesc(), offsets,
+ memDesc, offsets,
layout.dropSgLayoutAndData());
newOps.push_back(newOp);
}
@@ -924,19 +978,69 @@ struct WgToSgStoreMatrixOp : public OpConversionPattern<xegpu::StoreMatrixOp> {
matchAndRewrite(xegpu::StoreMatrixOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
+ xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
+
+ bool isPrivate = isPrivateMemDesc(op.getMemDesc());
+ auto memDesc = cast<TypedValue<xegpu::MemDescType>>(
+ isPrivate ? adaptor.getMemDesc()[0] : op.getMemDesc());
SmallVector<SmallVector<OpFoldResult>> offsetsList;
- if (failed(genOffsetsList(rewriter, op, offsetsList)))
+ if (isPrivate) {
+ offsetsList = genPrivateOffsetsList(rewriter, op.getLoc(),
+ op.getDataShape(), layout);
+ } else if (failed(genOffsetsList(rewriter, op, offsetsList))) {
return failure();
+ }
- xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
for (auto [v, offsets] : llvm::zip(adaptor.getData(), offsetsList))
- xegpu::StoreMatrixOp::create(rewriter, op.getLoc(), v, op.getMemDesc(),
- offsets, layout.dropSgLayoutAndData());
+ xegpu::StoreMatrixOp::create(rewriter, op.getLoc(), v, memDesc, offsets,
+ layout.dropSgLayoutAndData());
rewriter.eraseOp(op);
return success();
}
};
+// Shrinks a privatized create_mem_desc (marked by the pre-phase) to the
+// per-subgroup size: it allocates a smaller private buffer and builds a
+// mem_desc of the distributed shape, so each subgroup owns a private copy.
+struct WgToSgCreateMemDescOp
+ : public OpConversionPattern<xegpu::CreateMemDescOp> {
+ using OpConversionPattern<xegpu::CreateMemDescOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::CreateMemDescOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto layout = op->getAttrOfType<xegpu::DistributeLayoutAttr>(
+ kPrivatizeLayoutAttrName);
+ if (!layout)
+ return failure();
+
+ Location loc = op.getLoc();
+ xegpu::MemDescType mdescTy = op.getMemDesc().getType();
+ Type elemTy = mdescTy.getElementType();
+
+ // Per-subgroup distributed shape: the workgroup shape divided by sg_layout,
+ // holding every distribution round this subgroup owns (sg_data * rounds).
+ auto distShape =
+ layout.computeDistributedShape(llvm::to_vector(mdescTy.getShape()));
+ if (failed(distShape))
+ return failure();
+ SmallVector<int64_t> sgShape = *distShape;
+
+ // Allocate a private buffer just large enough for one subgroup's copy.
+ auto bytesPerElement = elemTy.getIntOrFloatBitWidth() / 8;
+ auto slmSize = computeProduct(sgShape) * bytesPerElement;
+ auto memrefTy =
+ MemRefType::get({slmSize}, rewriter.getI8Type(), {}, /*space=*/4);
+ auto buffer = memref::AllocaOp::create(rewriter, loc, memrefTy);
+
+ auto newMdescTy = xegpu::MemDescType::get(rewriter.getContext(), sgShape,
+ elemTy, mdescTy.getMemLayout());
+ auto newOp =
+ xegpu::CreateMemDescOp::create(rewriter, loc, newMdescTy, buffer);
+ rewriter.replaceOp(op, newOp.getResult());
+ return success();
+ }
+};
+
// This pattern distributes the vector.step ops to work at subgroup level
struct WgToSgVectorStepOp : public OpConversionPattern<vector::StepOp> {
using OpConversionPattern<vector::StepOp>::OpConversionPattern;
@@ -1529,7 +1633,8 @@ void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
WgToSgDpasMxOp, WgToSgPrefetchNdOp, WgToSgElementwiseOp,
WgToSgVectorBroadcastOp, WgToSgConvertLayoutOp,
WgToSgArithConstantOp, WgToSgLoadGatherOp, WgToSgStoreScatterOp,
- WgToSgLoadMatrixOp, WgToSgStoreMatrixOp, WgToSgVectorStepOp,
+ WgToSgLoadMatrixOp, WgToSgStoreMatrixOp, WgToSgCreateMemDescOp,
+ WgToSgVectorStepOp,
WgToSgVectorShapeCastOp, WgToSgMultiDimReductionOp,
WgToSgVectorTransposeOp, WgToSgVectorConstantMaskOp,
WgToSgVectorCreateMaskOp, WgToSgVectorBitCastOp,
@@ -1540,6 +1645,166 @@ void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
} // namespace mlir
namespace {
+
+//===----------------------------------------------------------------------===//
+// SLM privatization
+//===----------------------------------------------------------------------===//
+// An SLM buffer whose every matrix access (load_matrix/store_matrix) uses the
+// same subgroup layout, the same data (vector) shape, and the same offsets is
+// accessed identically by every subgroup. Because the same mem_desc implies the
+// same physical layout, an identical (sg_layout, sg_data, offset) means every
+// subgroup reads and writes the exact same SLM region, so the region is
+// effectively private to each subgroup and can be demoted to private memory
+// (space 4, backed by registers).
+//
+// This runs as a pre-phase of WG-to-SG distribution. For a qualifying buffer it
+// (a) flips the source memref's memory space from 3 (SLM) to 4 (private), and
+// (b) stashes the common workgroup layout on the create_mem_desc op as a
+// discardable marker attribute so the WgToSgCreateMemDescOp pattern can later
+// shrink the memref/mem_desc to the per-subgroup size. The matrix ops are left
+// intact for the regular distribution patterns that follow.
+
+// Compares two offset lists for structural equality. Static offsets compare by
+// value; dynamic offsets compare by SSA value identity.
+static bool sameOffsets(ArrayRef<OpFoldResult> a, ArrayRef<OpFoldResult> b) {
+ if (a.size() != b.size())
+ return false;
+ for (auto [lhs, rhs] : llvm::zip(a, b)) {
+ auto lhsAttr = dyn_cast<Attribute>(lhs);
+ auto rhsAttr = dyn_cast<Attribute>(rhs);
+ if (static_cast<bool>(lhsAttr) != static_cast<bool>(rhsAttr))
+ return false;
+ if (lhsAttr) {
+ if (lhsAttr != rhsAttr)
+ return false;
+ } else if (cast<Value>(lhs) != cast<Value>(rhs)) {
+ return false;
+ }
+ }
+ return true;
+}
+
+// Returns the single create_mem_desc op that views `source` if `source` is a
+// privatizable local buffer. The buffer must be a local alloca/alloc in SLM
+// whose only user is exactly one create_mem_desc op; any other user (or a
+// second view) means the buffer might escape or be aliased, so it is rejected.
+static xegpu::CreateMemDescOp getPrivatizableView(Value source) {
+ Operation *defOp = source.getDefiningOp();
+ if (!defOp || !isa<memref::AllocaOp, memref::AllocOp>(defOp))
+ return nullptr;
+
+ auto memrefTy = dyn_cast<MemRefType>(source.getType());
+ if (!memrefTy || !xegpu::XeGPUDialect::isSharedMemory(memrefTy))
+ return nullptr;
+
+ // The buffer must be viewed by exactly one create_mem_desc and nothing else,
+ // so it neither escapes nor aliases another mem_desc.
+ if (!source.hasOneUse())
+ return nullptr;
+ return dyn_cast<xegpu::CreateMemDescOp>(*source.getUsers().begin());
+}
+
+// If every load_matrix/store_matrix on `createOp`'s mem_desc is a workgroup
+// access sharing the same sg_layout, sg_data, data shape, and offsets, and the
+// workgroup tile is evenly distributed to the subgroups (no broadcast/overlap),
+// returns that common layout. Otherwise returns nullptr.
+static xegpu::DistributeLayoutAttr
+getSubgroupPrivateLayout(xegpu::CreateMemDescOp createOp) {
+ xegpu::DistributeLayoutAttr commonLayout;
+ SmallVector<OpFoldResult> commonOffsets;
+ ArrayRef<int64_t> commonDataShape;
+ bool seen = false;
+
+ auto checkAccess = [&](xegpu::DistributeLayoutAttr layout,
+ SmallVector<OpFoldResult> offsets,
+ ArrayRef<int64_t> dataShape) -> bool {
+ // Every access must be a distributed workgroup-level access.
+ if (!layout || !layout.isForWorkgroup())
+ return false;
+
+ // The workgroup tile must be evenly distributed across the subgroups with
+ // no broadcast: sg_layout[d] * sg_data[d] must not exceed the tile in any
+ // dim. A larger product means a dimension wraps around and is broadcast to
+ // multiple subgroups, so their regions overlap and are not private.
+ SmallVector<int64_t> sgLayout = layout.getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = layout.getEffectiveSgDataAsInt();
+ if (sgLayout.size() != dataShape.size() ||
+ sgData.size() != dataShape.size())
+ return false;
+ for (auto [l, d, tile] : llvm::zip_equal(sgLayout, sgData, dataShape))
+ if (l * d > tile)
+ return false;
+
+ if (!seen) {
+ commonLayout = layout;
+ commonOffsets = std::move(offsets);
+ commonDataShape = dataShape;
+ seen = true;
+ return true;
+ }
+ return commonLayout.getEffectiveSgLayoutAsInt() == sgLayout &&
+ commonLayout.getEffectiveSgDataAsInt() == sgData &&
+ commonDataShape == dataShape && sameOffsets(commonOffsets, offsets);
+ };
+
+ for (Operation *memUser : createOp.getMemDesc().getUsers()) {
+ if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(memUser)) {
+ if (!checkAccess(loadOp.getLayoutAttr(), loadOp.getMixedOffsets(),
+ loadOp.getDataShape()))
+ return nullptr;
+ } else if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(memUser)) {
+ if (!checkAccess(storeOp.getLayoutAttr(), storeOp.getMixedOffsets(),
+ storeOp.getDataShape()))
+ return nullptr;
+ } else {
+ // Any other use of the mem_desc is unexpected; be conservative.
+ return nullptr;
+ }
+ }
+
+ // Require at least one access to have been observed.
+ return seen ? commonLayout : nullptr;
+}
+
+// Rewrites the buffer defined by `defOp` (an alloca/alloc) so that its memory
+// space is private (4) instead of shared (3), updates the source type of the
+// create_mem_desc that consumes it, and marks that op with the common workgroup
+// layout so the conversion pattern can shrink it to the per-subgroup size.
+static void privatizeSlmBuffer(Operation *defOp, xegpu::CreateMemDescOp view,
+ xegpu::DistributeLayoutAttr layout) {
+ OpBuilder builder(defOp);
+ Value oldBuffer = defOp->getResult(0);
+ auto oldTy = cast<MemRefType>(oldBuffer.getType());
+ auto privateSpace = builder.getI64IntegerAttr(4);
+ auto newTy = MemRefType::get(oldTy.getShape(), oldTy.getElementType(),
+ oldTy.getLayout(), privateSpace);
+
+ Operation *newDefOp = builder.clone(*defOp);
+ newDefOp->getResult(0).setType(newTy);
+ oldBuffer.replaceAllUsesWith(newDefOp->getResult(0));
+ defOp->erase();
+
+ view->setAttr(kPrivatizeLayoutAttrName, layout);
+}
+
+// Entry point for the SLM privatization pre-phase.
+static void privatizeSharedLocalMemory(Operation *root) {
+ SmallVector<std::tuple<Operation *, xegpu::CreateMemDescOp,
+ xegpu::DistributeLayoutAttr>>
+ toPrivatize;
+
+ root->walk([&](xegpu::CreateMemDescOp createOp) {
+ Value source = createOp.getSource();
+ if (getPrivatizableView(source) != createOp)
+ return;
+ if (auto layout = getSubgroupPrivateLayout(createOp))
+ toPrivatize.push_back({source.getDefiningOp(), createOp, layout});
+ });
+
+ for (auto [defOp, view, layout] : toPrivatize)
+ privatizeSlmBuffer(defOp, view, layout);
+}
+
struct XeGPUWgToSgDistributePass
: public xegpu::impl::XeGPUWgToSgDistributeBase<XeGPUWgToSgDistributePass> {
void runOnOperation() override;
@@ -1549,11 +1814,19 @@ struct XeGPUWgToSgDistributePass
void XeGPUWgToSgDistributePass::runOnOperation() {
Operation *op = getOperation();
+
if (!xegpu::recoverTemporaryLayouts(op)) {
signalPassFailure();
return;
}
+ // Pre-phase: demote SLM buffers that are accessed identically by every
+ // subgroup (same sg_layout, sg_data, data shape and offsets) to subgroup-
+ // private memory. Run after recoverTemporaryLayouts, which strips discardable
+ // DistributeLayoutAttr-valued attrs (and would otherwise remove the marker
+ // attribute this phase sets on create_mem_desc).
+ privatizeSharedLocalMemory(op);
+
// Collect existing UnrealizedConversionCastOps. These must be preserved.
llvm::SmallSetVector<UnrealizedConversionCastOp, 8> existingCasts;
getOperation()->walk(
@@ -1625,6 +1898,13 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(op.getLayoutAttr());
});
+ // A create_mem_desc marked by the privatization pre-phase must be shrunk to
+ // the per-subgroup size; all others are already legal.
+ target.addDynamicallyLegalOp<xegpu::CreateMemDescOp>(
+ [=](xegpu::CreateMemDescOp op) -> bool {
+ return !op->hasAttr(kPrivatizeLayoutAttrName);
+ });
+
target.addDynamicallyLegalOp<arith::ConstantOp>(
[=](arith::ConstantOp op) -> bool {
auto vecType = dyn_cast<VectorType>(op.getType());
@@ -1703,4 +1983,12 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
// Fold cancelling cast chains and erase dead casts.
xegpu::cleanupUnrealizedConversionCasts(getOperation(), existingCasts);
xegpu::removeTemporaryLayoutAttrs(getOperation());
+
+ // Erase the original full-size buffers left dead by privatization: the
+ // WgToSgCreateMemDescOp pattern allocates fresh per-subgroup buffers, so the
+ // pre-phase's space-4 allocas no longer have any users.
+ getOperation()->walk([](Operation *op) {
+ if (isa<memref::AllocaOp, memref::AllocOp>(op) && op->use_empty())
+ op->erase();
+ });
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index d5d4950fe7d7e..157a3bcda7b1d 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -607,7 +607,7 @@ func.func @slice_attr_repeat_dim() {
// -----
func.func @create_mem_desc_non_slm() {
%m = memref.alloca() {alignment = 1024} : memref<2048xi8, 1>
- // expected-error at +1 {{operand #0 must be reside in share memory and statically 1d shaped memref }}
+ // expected-error at +1 {{operand #0 must be reside in shared or private memory and statically 1d shaped memref }}
%mem_desc = xegpu.create_mem_desc %m : memref<2048xi8, 1> -> !xegpu.mem_desc<16x64xf16>
return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-privatize.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-privatize.mlir
new file mode 100644
index 0000000000000..57f2a01189cee
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-privatize.mlir
@@ -0,0 +1,207 @@
+// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s
+
+// The SLM privatization pre-phase of WG-to-SG distribution demotes a shared
+// local memory (space 3) buffer to subgroup-private memory (space 4) when every
+// matrix access to it uses the same sg_layout, sg_data, data shape and offsets,
+// and the workgroup tile is evenly distributed to the subgroups (no broadcast).
+// In that case each subgroup reads/writes an identical, non-overlapping region,
+// so the buffer and mem_desc are shrunk to the per-subgroup size and the matrix
+// ops are re-indexed with local, subgroup-id-free offsets.
+
+gpu.module @test {
+ // A buffer whose load_matrix and store_matrix share the same sg_layout,
+ // sg_data and offsets is private to each subgroup: it is moved to memory space
+ // 4 and shrunk from 64x128 to the per-subgroup 32x32 tile.
+ // CHECK-LABEL: gpu.func @privatize_same_offset_layout
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<4096xi8, 4>
+ // CHECK: %[[MD:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<4096xi8, 4> -> !xegpu.mem_desc<32x32xf32>
+ // CHECK: xegpu.load_matrix %[[MD]][0, 0] : !xegpu.mem_desc<32x32xf32> -> vector<32x32xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD]][0, 0] : vector<32x32xf32>, !xegpu.mem_desc<32x32xf32>
+ gpu.func @privatize_same_offset_layout() {
+ %cst = arith.constant dense<1.0> : vector<64x128xf32>
+ %a = memref.alloca() : memref<32768xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ %ld = xegpu.load_matrix %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : !xegpu.mem_desc<64x128xf32> -> vector<64x128xf32>
+ xegpu.store_matrix %cst, %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // When a dimension spans multiple distribution rounds (128 / sg_layout 8 = 16
+ // = 2 rounds of sg_data 8), the shrunk buffer holds both rounds and each round
+ // is indexed with a local offset (0 and 8).
+ // CHECK-LABEL: gpu.func @privatize_multi_round
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<64xi8, 4>
+ // CHECK: %[[MD:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<64xi8, 4> -> !xegpu.mem_desc<16xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD]][0] : vector<8xf32>, !xegpu.mem_desc<16xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD]][8] : vector<8xf32>, !xegpu.mem_desc<16xf32>
+ // CHECK: xegpu.load_matrix %[[MD]][0] : !xegpu.mem_desc<16xf32> -> vector<8xf32>
+ // CHECK: xegpu.load_matrix %[[MD]][8] : !xegpu.mem_desc<16xf32> -> vector<8xf32>
+ gpu.func @privatize_multi_round() {
+ %cst = arith.constant dense<1.0> : vector<128xf32>
+ %a = memref.alloca() : memref<512xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<512xi8, 3> -> !xegpu.mem_desc<128xf32>
+ xegpu.store_matrix %cst, %md[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : vector<128xf32>, !xegpu.mem_desc<128xf32>
+ %l = xegpu.load_matrix %md[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : !xegpu.mem_desc<128xf32> -> vector<128xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // Accesses use different offsets, so the region seen by each subgroup is not
+ // identical: the buffer stays in shared local memory (space 3), full size.
+ // CHECK-LABEL: gpu.func @no_privatize_diff_offsets
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<65536xi8, 3>
+ // CHECK: xegpu.create_mem_desc %[[ALLOCA]] : memref<65536xi8, 3> -> !xegpu.mem_desc<128x128xf32>
+ gpu.func @no_privatize_diff_offsets() {
+ %cst = arith.constant dense<1.0> : vector<64x128xf32>
+ %a = memref.alloca() : memref<65536xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<65536xi8, 3> -> !xegpu.mem_desc<128x128xf32>
+ %ld = xegpu.load_matrix %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : !xegpu.mem_desc<128x128xf32> -> vector<64x128xf32>
+ xegpu.store_matrix %cst, %md[32, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : vector<64x128xf32>, !xegpu.mem_desc<128x128xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // Accesses use different sg_layouts, so subgroups partition the region
+ // differently: the buffer stays in shared local memory (space 3).
+ // CHECK-LABEL: gpu.func @no_privatize_diff_layout
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<32768xi8, 3>
+ // CHECK: xegpu.create_mem_desc %[[ALLOCA]] : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ gpu.func @no_privatize_diff_layout() {
+ %cst = arith.constant dense<1.0> : vector<64x128xf32>
+ %a = memref.alloca() : memref<32768xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ %ld = xegpu.load_matrix %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : !xegpu.mem_desc<64x128xf32> -> vector<64x128xf32>
+ xegpu.store_matrix %cst, %md[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 64]>}>
+ : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // Accesses use different data (vector) shapes, so the per-subgroup regions
+ // differ: the buffer stays in shared local memory (space 3).
+ // CHECK-LABEL: gpu.func @no_privatize_diff_data_shape
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<32768xi8, 3>
+ // CHECK: xegpu.create_mem_desc %[[ALLOCA]] : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ gpu.func @no_privatize_diff_data_shape() {
+ %c0 = arith.constant dense<1.0> : vector<64x128xf32>
+ %c1 = arith.constant dense<1.0> : vector<32x128xf32>
+ %a = memref.alloca() : memref<32768xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ xegpu.store_matrix %c0, %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
+ %ld = xegpu.load_matrix %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 32]>}>
+ : !xegpu.mem_desc<64x128xf32> -> vector<32x128xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // The workgroup tile wraps around (sg_layout 4 * sg_data 8 = 32 > tile 8), so
+ // the single tile is broadcast to all four subgroups and their regions
+ // overlap: the buffer is not private and stays in shared local memory.
+ // CHECK-LABEL: gpu.func @no_privatize_broadcast
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<32xi8, 3>
+ // CHECK: xegpu.create_mem_desc %[[ALLOCA]] : memref<32xi8, 3> -> !xegpu.mem_desc<8xf32>
+ gpu.func @no_privatize_broadcast() {
+ %cst = arith.constant dense<1.0> : vector<8xf32>
+ %a = memref.alloca() : memref<32xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<32xi8, 3> -> !xegpu.mem_desc<8xf32>
+ xegpu.store_matrix %cst, %md[0] <{layout = #xegpu.layout<sg_layout = [4], sg_data = [8]>}>
+ : vector<8xf32>, !xegpu.mem_desc<8xf32>
+ %l = xegpu.load_matrix %md[0] <{layout = #xegpu.layout<sg_layout = [4], sg_data = [8]>}>
+ : !xegpu.mem_desc<8xf32> -> vector<8xf32>
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // The buffer escapes the function (passed to a call), so it may be observed
+ // by other subgroups and must not be privatized: it stays in space 3.
+ func.func private @use(memref<32768xi8, 3>)
+ // CHECK-LABEL: gpu.func @no_privatize_escaping_buffer
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<32768xi8, 3>
+ // CHECK: xegpu.create_mem_desc %[[ALLOCA]] : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ gpu.func @no_privatize_escaping_buffer() {
+ %cst = arith.constant dense<1.0> : vector<64x128xf32>
+ %a = memref.alloca() : memref<32768xi8, 3>
+ %md = xegpu.create_mem_desc %a : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
+ xegpu.store_matrix %cst, %md[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}>
+ : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
+ func.call @use(%a) : (memref<32768xi8, 3>) -> ()
+ gpu.return
+ }
+}
+
+// -----
+
+gpu.module @test {
+ // End-to-end softmax-like kernel (after vector-to-xegpu, rewritten to use
+ // load_matrix/store_matrix). Both SLM scratch buffers (running max and running
+ // sum) are read and written by every subgroup with the same sg_layout, sg_data
+ // and offset, so both are demoted to subgroup-private memory and shrunk from
+ // 64 elements to the per-subgroup 8 elements.
+ // CHECK-LABEL: gpu.func @payload_kernel
+ // CHECK-DAG: %[[MAX:.*]] = memref.alloca() : memref<32xi8, 4>
+ // CHECK-DAG: %[[MD_MAX:.*]] = xegpu.create_mem_desc %[[MAX]] : memref<32xi8, 4> -> !xegpu.mem_desc<8xf32>
+ // CHECK-DAG: %[[SUM:.*]] = memref.alloca() : memref<32xi8, 4>
+ // CHECK-DAG: %[[MD_SUM:.*]] = xegpu.create_mem_desc %[[SUM]] : memref<32xi8, 4> -> !xegpu.mem_desc<8xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD_MAX]][0] : vector<8xf32>, !xegpu.mem_desc<8xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD_SUM]][0] : vector<8xf32>, !xegpu.mem_desc<8xf32>
+ // CHECK: scf.for
+ // CHECK: xegpu.load_matrix %[[MD_MAX]][0] : !xegpu.mem_desc<8xf32> -> vector<8xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD_MAX]][0] : vector<8xf32>, !xegpu.mem_desc<8xf32>
+ // CHECK: xegpu.load_matrix %[[MD_SUM]][0] : !xegpu.mem_desc<8xf32> -> vector<8xf32>
+ // CHECK: xegpu.store_matrix %{{.*}}, %[[MD_SUM]][0] : vector<8xf32>, !xegpu.mem_desc<8xf32>
+ gpu.func @payload_kernel() kernel {
+ %c16 = arith.constant 16 : index
+ %c512 = arith.constant 512 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<64xf32>
+ %c0 = arith.constant 0 : index
+ %cst_0 = arith.constant dense<0xFFC00000> : vector<64xf32>
+ %alloca = memref.alloca() : memref<256xi8, 3>
+ %alloca_1 = memref.alloca() : memref<256xi8, 3>
+ %md_max = xegpu.create_mem_desc %alloca_1 : memref<256xi8, 3> -> !xegpu.mem_desc<64xf32>
+ %md_sum = xegpu.create_mem_desc %alloca : memref<256xi8, 3> -> !xegpu.mem_desc<64xf32>
+ xegpu.store_matrix %cst_0, %md_max[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : vector<64xf32>, !xegpu.mem_desc<64xf32>
+ xegpu.store_matrix %cst, %md_sum[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : vector<64xf32>, !xegpu.mem_desc<64xf32>
+ scf.for %arg2 = %c0 to %c512 step %c16 {
+ %4 = xegpu.load_matrix %md_max[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : !xegpu.mem_desc<64xf32> -> vector<64xf32>
+ %5 = arith.subf %cst, %4 : vector<64xf32>
+ xegpu.store_matrix %5, %md_max[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : vector<64xf32>, !xegpu.mem_desc<64xf32>
+ %11 = xegpu.load_matrix %md_sum[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : !xegpu.mem_desc<64xf32> -> vector<64xf32>
+ %12 = arith.addf %11, %5 : vector<64xf32>
+ xegpu.store_matrix %12, %md_sum[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [8]>}>
+ : vector<64xf32>, !xegpu.mem_desc<64xf32>
+ }
+ gpu.return
+ }
+}
More information about the Mlir-commits
mailing list