[Mlir-commits] [mlir] [MLIR][XeGPU] Add unroll pass for XeGPU (PR #137010)

[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 HEAD~1 HEAD --extensions h,cpp -- mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp mlir/include/mlir/Dialect/XeGPU/IR/XeGPU.h mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
``````````

</details>

<details>
<summary>
View the diff from clang-format here.
</summary>

``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 94a236cb1..3d88fecec 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -210,7 +210,8 @@ struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
         mixedOffsets[x] = addi(oldX, subOffX);
         mixedOffsets[y] = addi(oldY, subOffY);
         auto newOp = rewriter.create<xegpu::CreateNdDescOp>(
-          loc, newTdescTy, op.getSource(), mixedOffsets, op.getMixedSizes(), op.getMixedStrides());
+            loc, newTdescTy, op.getSource(), mixedOffsets, op.getMixedSizes(),
+            op.getMixedStrides());
         newOps.push_back(newOp);
       }
     }
@@ -262,19 +263,22 @@ struct UnrollStoreNdOp : public UnrollPattern<xegpu::StoreNdOp> {
 
     auto elemTy = tdescTy.getElementType();
     auto newValueTy = valueTy.cloneWith(targetShape, elemTy);
-    auto newTdescTy = xegpu::TensorDescType::get(ctx, targetShape, elemTy, tdescTy.getEncoding(),
-        getLaneLayoutAttr(layout));
+    auto newTdescTy = xegpu::TensorDescType::get(ctx, targetShape, elemTy,
+                                                 tdescTy.getEncoding(),
+                                                 getLaneLayoutAttr(layout));
 
-    auto numNewOps = std::accumulate(grids.begin(), grids.end(), 1, std::multiplies<int64_t>());
+    auto numNewOps = std::accumulate(grids.begin(), grids.end(), 1,
+                                     std::multiplies<int64_t>());
     llvm::SmallVector<Type> convertedValTypes(numNewOps, newValueTy);
     llvm::SmallVector<Type> convertedTileTypes(numNewOps, newTdescTy);
-    auto convertedValues = addPackOp(op.getValue(), convertedValTypes, targetShape, loc, rewriter);
-    auto convertedTdescs = addPackOp(op.getTensorDesc(), convertedTileTypes, targetShape, loc, rewriter);
+    auto convertedValues =
+        addPackOp(op.getValue(), convertedValTypes, targetShape, loc, rewriter);
+    auto convertedTdescs = addPackOp(op.getTensorDesc(), convertedTileTypes,
+                                     targetShape, loc, rewriter);
 
     for (auto [v, t] : llvm::zip(convertedValues, convertedTdescs)) {
       rewriter.create<xegpu::StoreNdOp>(loc, v, t, op.getL1HintAttr(),
-                                           op.getL2HintAttr(),
-                                           op.getL3HintAttr());
+                                        op.getL2HintAttr(), op.getL3HintAttr());
     }
     rewriter.eraseOp(op);
     return success();
@@ -356,31 +360,32 @@ struct XeGPUUnrollPass final
 
   void runOnOperation() override {
     vector::UnrollVectorOptions options;
-    options.setNativeShapeFn(
-        [&](Operation *op) -> std::optional<SmallVector<int64_t>> {
-          if (isa<xegpu::CreateNdDescOp, xegpu::LoadNdOp, xegpu::StoreNdOp>(op)) {
-            xegpu::TensorDescType tdescTy;
-            if (auto createNdOp = dyn_cast<xegpu::CreateNdDescOp>(op)) {
-              tdescTy = createNdOp.getType();
-            } else if (auto loadNdOp = dyn_cast<xegpu::LoadNdOp>(op)) {
-              tdescTy = loadNdOp.getTensorDescType();
-            } else if (auto storeNdOp = dyn_cast<xegpu::StoreNdOp>(op)) {
-              tdescTy = storeNdOp.getTensorDescType();
-            }
-
-            if (auto layout = tdescTy.getLayoutAttr()) {
-              if (auto inst_data = layout.getInstData())
-                return SmallVector<int64_t>(inst_data.asArrayRef().begin(),
-                                            inst_data.asArrayRef().end());
-            }
-          }
-
-          return std::nullopt;
-        });
+    options.setNativeShapeFn([&](Operation *op)
+                                 -> std::optional<SmallVector<int64_t>> {
+      if (isa<xegpu::CreateNdDescOp, xegpu::LoadNdOp, xegpu::StoreNdOp>(op)) {
+        xegpu::TensorDescType tdescTy;
+        if (auto createNdOp = dyn_cast<xegpu::CreateNdDescOp>(op)) {
+          tdescTy = createNdOp.getType();
+        } else if (auto loadNdOp = dyn_cast<xegpu::LoadNdOp>(op)) {
+          tdescTy = loadNdOp.getTensorDescType();
+        } else if (auto storeNdOp = dyn_cast<xegpu::StoreNdOp>(op)) {
+          tdescTy = storeNdOp.getTensorDescType();
+        }
+
+        if (auto layout = tdescTy.getLayoutAttr()) {
+          if (auto inst_data = layout.getInstData())
+            return SmallVector<int64_t>(inst_data.asArrayRef().begin(),
+                                        inst_data.asArrayRef().end());
+        }
+      }
+
+      return std::nullopt;
+    });
 
     auto funcOp = getOperation();
     RewritePatternSet patterns(&getContext());
-    patterns.add<UnrollCreateNdOp, UnrollStoreNdOp>(patterns.getContext(), options);
+    patterns.add<UnrollCreateNdOp, UnrollStoreNdOp>(patterns.getContext(),
+                                                    options);
 
     // GreedyRewriteConfig config;
     // config.fold = false;

``````````

</details>


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