[Mlir-commits] [mlir] [SCF] add debug label for MergeNestedParallelLoops pattern (PR #135980)

[Scott Manley]

https://github.com/rscottmanley updated https://github.com/llvm/llvm-project/pull/135980

>From d0974e17c5588c7f58e7a5828bc14eafe8fa07af Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Wed, 16 Apr 2025 08:56:09 -0700
Subject: [PATCH 1/2] [SCF] add debug label for MergeNestedParallelLoops
 pattern

Adding debug labels to rewrite patterns gives the canonicalizer the
ability to disable specific patterns if undesired.

The motivation here is that since there is currently no way to handle
attributes on ops that are changed when these patterns are applied, this
at least provides a workaround.
---
 mlir/lib/Dialect/SCF/IR/SCF.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index c454f342b4204..04e89105c056a 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3202,9 +3202,9 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
 
 void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results
-      .add<ParallelOpSingleOrZeroIterationDimsFolder, MergeNestedParallelLoops>(
-          context);
+  results.add<ParallelOpSingleOrZeroIterationDimsFolder>(context);
+  results.addWithLabel<MergeNestedParallelLoops>({"MergeNestedParallelLoops"},
+                                                 context);
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,

>From b4275ee08ad335f0762bd1ffd2c0e3f8ba4d516c Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Wed, 16 Apr 2025 12:47:11 -0700
Subject: [PATCH 2/2] change to setDebugName for pattern only

---
 mlir/lib/Dialect/SCF/IR/SCF.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 04e89105c056a..4aab0608b47e0 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3143,6 +3143,8 @@ struct ParallelOpSingleOrZeroIterationDimsFolder
 struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
   using OpRewritePattern<ParallelOp>::OpRewritePattern;
 
+  void initialize() { setDebugName("MergeNestedParallelLoops"); }
+
   LogicalResult matchAndRewrite(ParallelOp op,
                                 PatternRewriter &rewriter) const override {
     Block &outerBody = *op.getBody();
@@ -3202,9 +3204,9 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
 
 void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results.add<ParallelOpSingleOrZeroIterationDimsFolder>(context);
-  results.addWithLabel<MergeNestedParallelLoops>({"MergeNestedParallelLoops"},
-                                                 context);
+  results
+      .add<ParallelOpSingleOrZeroIterationDimsFolder, MergeNestedParallelLoops>(
+          context);
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,