[Mlir-commits] [mlir] [MLIR][XeGPU] Refactor XeGPU layout propagation: passing lane_layout/lane_data with inst_data (PR #202868)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 9 23:49:02 PDT 2026
github-actions[bot] wrote:
<!--LLVM CODE FORMAT COMMENT: {clang-format}-->
:warning: C/C++ code formatter, clang-format found issues in your code. :warning:
<details>
<summary>
You can test this locally with the following command:
</summary>
``````````bash
git-clang-format --diff origin/main HEAD --extensions cpp,h -- mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp --diff_from_common_commit
``````````
:warning:
The reproduction instructions above might return results for more than one PR
in a stack if you are using a stacked PR workflow. You can limit the results by
changing `origin/main` to the base branch/commit you want to compare against.
:warning:
</details>
<details>
<summary>
View the diff from clang-format here.
</summary>
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index f3cbf6ae7..44b5441fb 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -245,8 +245,7 @@ completeLoadGatherLayoutFromInstData(DistributeLayoutAttr consumerLayout,
DistributeLayoutAttr
completeStoreScatterLayoutFromInstData(DistributeLayoutAttr consumerLayout,
- Type elemTy,
- const uArch::uArch *uArch);
+ Type elemTy, const uArch::uArch *uArch);
/// Sets up the anchor layout for a store_nd operation. StoreNd does not
/// consider a consumer layout (it is a data sink), and picks its layout from
@@ -269,11 +268,10 @@ DistributeLayoutAttr setupPrefetchNdAnchorLayout(LayoutKind layoutKind,
/// Otherwise defaults derived from uArch block parameters are used.
/// `consumerLayout` may be null. `numSg` is only used for Subgroup-kind
/// layouts when the consumer does not already provide an sg_layout.
-DistributeLayoutAttr setupLoadNdAnchorLayout(LayoutKind layoutKind,
- VectorType vectorTy,
- DistributeLayoutAttr consumerLayout,
- int numSg,
- const uArch::uArch *uArch);
+DistributeLayoutAttr
+setupLoadNdAnchorLayout(LayoutKind layoutKind, VectorType vectorTy,
+ DistributeLayoutAttr consumerLayout, int numSg,
+ const uArch::uArch *uArch);
/// Sets up the anchor layouts for a dpas operands (A, B, and C/D).
/// The numSg and consumerLayout (optional) are only used by sg layout creation.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 5cbb5cbc7..a170aca48 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -845,12 +845,10 @@ static bool leadingDimsAreUnit(ArrayRef<int64_t> shape, int numInnerDims) {
/// result layout can later be distributed without re-deriving the lane
/// layout. `instData`, `laneLayout`, and `laneData` may have different
/// element types; they are normalized to int32 entries.
-static xegpu::LayoutAttr
-buildInstDataLayoutWithLane(mlir::MLIRContext *context,
- ArrayRef<int64_t> instData,
- ArrayRef<int64_t> laneLayout,
- ArrayRef<int64_t> laneData,
- DenseI32ArrayAttr orderAttr = nullptr) {
+static xegpu::LayoutAttr buildInstDataLayoutWithLane(
+ mlir::MLIRContext *context, ArrayRef<int64_t> instData,
+ ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
+ DenseI32ArrayAttr orderAttr = nullptr) {
auto toI32Attr = [&](auto range) {
SmallVector<int32_t> v(range.begin(), range.end());
return DenseI32ArrayAttr::get(context, v);
@@ -1714,8 +1712,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
std::min<int64_t>(uArch->getSubgroupSize(), innermostShape);
laneLayout.back() = lanesOnInnermost;
if (lanesOnInnermost > 0)
- laneData.back() = std::min<int64_t>(packingFactor,
- innermostShape / lanesOnInnermost);
+ laneData.back() =
+ std::min<int64_t>(packingFactor, innermostShape / lanesOnInnermost);
else
laneData.back() = 1;
@@ -1765,8 +1763,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
// (e.g. sub-byte floats with no uArch entry), fall back to
// lane_layout * lane_data (k = 1).
SmallVector<int64_t> instData(rank, 1);
- int instWidth = xegpu::getLargestDivisor(
- static_cast<int>(dataShape.back()), bWidths);
+ int instWidth =
+ xegpu::getLargestDivisor(static_cast<int>(dataShape.back()), bWidths);
if (instWidth == -1)
instData.back() = laneLayout.back() * laneData.back();
else
@@ -1814,9 +1812,8 @@ static xegpu::DistributeLayoutAttr setupGenericNdAnchorLayout(
return nullptr;
SmallVector<int> sgLayout = {static_cast<int>(sgLayouts[0].first),
static_cast<int>(sgLayouts[0].second)};
- SmallVector<int> sgData = {
- static_cast<int>(dataShape[0]) / sgLayout[0],
- static_cast<int>(dataShape[1]) / sgLayout[1]};
+ SmallVector<int> sgData = {static_cast<int>(dataShape[0]) / sgLayout[0],
+ static_cast<int>(dataShape[1]) / sgLayout[1]};
return xegpu::LayoutAttr::get(
context, DenseI32ArrayAttr::get(context, sgLayout),
DenseI32ArrayAttr::get(context, sgData),
@@ -1886,10 +1883,11 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
/// against uArch constraints; if valid, the consumer's `inst_data` /
/// `sg_layout` are honored. Otherwise the helper falls back to defaults
/// derived from uArch block parameters.
-xegpu::DistributeLayoutAttr xegpu::setupLoadNdAnchorLayout(
- xegpu::LayoutKind layoutKind, VectorType resVecTy,
- xegpu::DistributeLayoutAttr consumerLayout, int numSg,
- const xegpu::uArch::uArch *uArch) {
+xegpu::DistributeLayoutAttr
+xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
+ VectorType resVecTy,
+ xegpu::DistributeLayoutAttr consumerLayout,
+ int numSg, const xegpu::uArch::uArch *uArch) {
auto context = resVecTy.getContext();
Type elemTy = resVecTy.getElementType();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index cb9d6eeed..a46415868 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -747,8 +747,8 @@ void LayoutInfoPropagation::visitShapeCastOp(
xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
llvm::dbgs() << "[DEBUG visitShapeCast] op=" << *shapeCast
- << " consumer=" << resultLayoutAttr
- << " src=" << srcLayoutAttr << "\n";
+ << " consumer=" << resultLayoutAttr << " src=" << srcLayoutAttr
+ << "\n";
propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(srcLayoutAttr)));
}
@@ -1332,8 +1332,7 @@ void LayoutInfoPropagation::visitStoreMatrixOp(
ArrayRef<const LayoutInfoLattice *> results) {
xegpu::DistributeLayoutAttr anchorLayout = storeMatrix.getLayoutAttr();
LayoutInfo layout;
- VectorType srcVecTy =
- llvm::cast<VectorType>(storeMatrix.getData().getType());
+ VectorType srcVecTy = llvm::cast<VectorType>(storeMatrix.getData().getType());
const uArch *uArch = getUArch(getChipStr(storeMatrix).value_or(""));
if (!uArch)
return;
``````````
</details>
https://github.com/llvm/llvm-project/pull/202868
More information about the Mlir-commits
mailing list