[mlir][affine] Remove one-element linearize_index as a canonicalization

[krzysz00]

By analogy to the canonicalization for affine.delinearize_index, remove affine.linearize_index ops that only have one multi-index input.

Example:

Canonicalize

%1 = affine.linearize_index [%0] by (64)

to

%1 = %0

While I'm here, get rid of an extra space in the syntax.

[llvmbot]

@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Krzysztof Drewniak (krzysz00)

Changes

By analogy to the canonicalization for affine.delinearize_index, remove affine.linearize_index ops that only have one multi-index input.

While I'm here, get rid of an extra space in the syntax.

---

Full diff: https://github.com/llvm/llvm-project/pull/115542.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+1-1)
• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+19-1)
• (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+11)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 92fd6e99338ae2..1dd9b9a440ecc8 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -1158,7 +1158,7 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
 
   let assemblyFormat = [{
     (`disjoint` $disjoint^)? ` `
-    `[` $multi_index `]` `by` ` `
+    `[` $multi_index `]` `by`
     custom<DynamicIndexList>($dynamic_basis, $static_basis, "::mlir::AsmParser::Delimiter::Paren")
     attr-dict `:` type($linear_index)
   }];
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index ca55c44856d191..3d38de4bf1068e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4734,11 +4734,29 @@ struct DropLinearizeUnitComponentsIfDisjointOrZero final
     return success();
   }
 };
+
+/// Rewrite `affine.linearize_index [%%x] by (%b)`, into `%x`.
+///
+/// By definition, that operation is `affine.apply affine_map<()[s0] -> (s0)>,`
+/// which is the identity.
+struct DropLinearizeOneBasisElement final
+    : OpRewritePattern<affine::AffineLinearizeIndexOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineLinearizeIndexOp op,
+                                PatternRewriter &rewriter) const override {
+    if (op.getStaticBasis().size() != 1 || op.getMultiIndex().size() != 1)
+      return failure();
+    rewriter.replaceOp(op, op.getMultiIndex().front());
+    return success();
+  }
+};
 } // namespace
 
 void affine::AffineLinearizeIndexOp::getCanonicalizationPatterns(
     RewritePatternSet &patterns, MLIRContext *context) {
-  patterns.add<DropLinearizeUnitComponentsIfDisjointOrZero>(context);
+  patterns.add<DropLinearizeUnitComponentsIfDisjointOrZero,
+               DropLinearizeOneBasisElement>(context);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index f6007aa16c1266..fa179744094c67 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1566,3 +1566,14 @@ func.func @linearize_all_zero_unit_basis() -> index {
   %ret = affine.linearize_index [%c0, %c0] by (1, 1) : index
   return %ret : index
 }
+
+// -----
+
+// CHECK-LABEL: @linearize_one_element_basis
+// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index)
+// CHECK-NOT: affine.linearize_index
+// CHECK: return %[[arg0]]
+func.func @linearize_one_element_basis(%arg0: index, %arg1: index) -> index {
+  %ret = affine.linearize_index [%arg0] by (%arg1) : index
+  return %ret : index
+}

[joker-eph]

By analogy to the canonicalization for affine.delinearize_index, remove affine.linearize_index ops that only have one multi-index input.

Can you elaborate on this and provide a more complete / standalone description. Thanks :)

[joker-eph]

By analogy to the canonicalization for affine.delinearize_index, remove affine.linearize_index ops that only have one multi-index input.

Can you elaborate on this and provide a more complete / standalone description. Thanks :)

Actually re-reading this 3 times, I think it makes sense, if you find a cleared way to express this that would be welcome, otherwise just ignore me here :)

[krzysz00]

I added an example to clarify