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

[Han-Chung Wang]

================
@@ -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();
+    }
+
+    if (!newOperands.empty()) {
+      auto newDelinearizeOp = rewriter.create<affine::AffineDelinearizeIndexOp>(
+          loc, delinearizeOp.getLinearIndex(), newOperands);
+      int newIndex = 0;
+      // Map back the new delinearized indices to the values they replace.
+      for (auto i : llvm::seq<size_t>(0, replacements.size())) {
----------------
hanhanW wrote:

can this just be `for (auto &replacement : replacements)`?

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