[Mlir-commits] [mlir] [mlir][xegpu] SLM Privatization for WG-to-SG Distribution (PR #208921)

[llvmlistbot at llvm.org]

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 -- mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.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/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index bc5407082..0c43c925b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -37,8 +37,9 @@ 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.
+// 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__";
 
@@ -913,10 +914,11 @@ genPrivateOffsetsList(ConversionPatternRewriter &rewriter, Location loc,
   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.
+  // 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];
@@ -945,14 +947,15 @@ struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {
     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.
+    // 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);
+      offsetsList =
+          genPrivateOffsetsList(rewriter, op.getLoc(), wgShape, layout);
     } else if (failed(genOffsetsList(rewriter, op, offsetsList))) {
       return failure();
     }
@@ -961,9 +964,9 @@ struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {
     VectorType newResTy = VectorType::get(sgShape, elemTy);
     SmallVector<Value> newOps;
     for (auto offsets : offsetsList) {
-      auto newOp = xegpu::LoadMatrixOp::create(rewriter, op.getLoc(), newResTy,
-                                               memDesc, offsets,
-                                               layout.dropSgLayoutAndData());
+      auto newOp =
+          xegpu::LoadMatrixOp::create(rewriter, op.getLoc(), newResTy, memDesc,
+                                      offsets, layout.dropSgLayoutAndData());
       newOps.push_back(newOp);
     }
     rewriter.replaceOpWithMultiple(op, {newOps});
@@ -1634,12 +1637,11 @@ void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
                WgToSgVectorBroadcastOp, WgToSgConvertLayoutOp,
                WgToSgArithConstantOp, WgToSgLoadGatherOp, WgToSgStoreScatterOp,
                WgToSgLoadMatrixOp, WgToSgStoreMatrixOp, WgToSgCreateMemDescOp,
-               WgToSgVectorStepOp,
-               WgToSgVectorShapeCastOp, WgToSgMultiDimReductionOp,
-               WgToSgVectorTransposeOp, WgToSgVectorConstantMaskOp,
-               WgToSgVectorCreateMaskOp, WgToSgVectorBitCastOp,
-               WgToSgVectorInterleaveOp, WgToSgVectorDeinterleaveOp>(
-      patterns.getContext());
+               WgToSgVectorStepOp, WgToSgVectorShapeCastOp,
+               WgToSgMultiDimReductionOp, WgToSgVectorTransposeOp,
+               WgToSgVectorConstantMaskOp, WgToSgVectorCreateMaskOp,
+               WgToSgVectorBitCastOp, WgToSgVectorInterleaveOp,
+               WgToSgVectorDeinterleaveOp>(patterns.getContext());
 }
 } // namespace xegpu
 } // namespace mlir

``````````

</details>


https://github.com/llvm/llvm-project/pull/208921