[Mlir-commits] [mlir] [mlir] Add forall canonicalization to replace constant induction vars (PR #112764)

[Han-Chung Wang]

================
@@ -1767,6 +1767,32 @@ struct ForallOpSingleOrZeroIterationDimsFolder
   }
 };
 
+struct ForallOpReplaceConstantInductionVar : public OpRewritePattern<ForallOp> {
+  using OpRewritePattern<ForallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ForallOp op,
+                                PatternRewriter &rewriter) const override {
+    // Replace all induction vars with a single trip count with their lower
+    // bound.
+    Location loc = op.getLoc();
+    bool replacedIv = false;
+    for (auto [lb, ub, step, iv] :
+         llvm::zip(op.getMixedLowerBound(), op.getMixedUpperBound(),
+                   op.getMixedStep(), op.getInductionVars())) {
+      if (iv.getUses().begin() == iv.getUses().end())
+        continue;
+      auto numIterations = constantTripCount(lb, ub, step);
+      if (!numIterations.has_value() || numIterations.value() != 1) {
+        continue;
+      }
+      rewriter.replaceAllUsesWith(
+          iv, getValueOrCreateConstantIndexOp(rewriter, loc, lb));
+      return success();
+    }
+    return failure();
----------------
hanhanW wrote:

I think we can make a slight improvement here which avoids calling the pattern many times. The `success()` method can take a boolean value. So I'd suggest to have a `bool changed=false;` variable and we return `success(changed)` at the end.

```suggestion
      changed = true;
    }
    return success(changed);
```

https://github.com/llvm/llvm-project/blob/e67442486d5efd48235f62b438557bc95193fc48/llvm/include/llvm/Support/LogicalResult.h#L53-L57

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