[Mlir-commits] [mlir] [mlir][SCF] Use Affine ops for indexing math. (PR #108450)

[llvmlistbot at llvm.org]

================
@@ -4534,6 +4534,140 @@ LogicalResult AffineDelinearizeIndexOp::verify() {
   return success();
 }
 
+namespace {
+
+// Drops delinearization indices that correspond to unit-extent basis
+struct DropUnitExtentBasis
+    : public OpRewritePattern<affine::AffineDelinearizeIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineDelinearizeIndexOp delinearizeOp,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> replacements(delinearizeOp->getNumResults(), nullptr);
+    std::optional<Value> zero = std::nullopt;
+    Location loc = delinearizeOp->getLoc();
+    auto getZero = [&]() -> Value {
+      if (!zero) {
+        zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+      }
+      return zero.value();
+    };
+
+    // Replace all indices corresponding to unit-extent basis with 0.
+    // Remaining basis can be used to get a new `affine.delinearize_index` op.
+    SmallVector<Value> newOperands;
+    for (auto [index, basis] : llvm::enumerate(delinearizeOp.getBasis())) {
+      if (matchPattern(basis, m_One())) {
+        replacements[index] = getZero();
+      } else {
+        newOperands.push_back(basis);
+      }
+    }
+
+    if (newOperands.size() == delinearizeOp.getBasis().size()) {
+      return failure();
+    }
----------------
MaheshRavishankar wrote:

I am trying to avoid creating IR if we dont need to. Seen it lead to some bad issues with pattern rewrites when we do this (especially with Affine ops involved).

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