[mlir][affine] Remove one-element linearize_index as a canonicalization (#115542)

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

Example:

Canonicalize
```mlir
%1 = affine.linearize_index [%0] by (64)
```
to
```mlir
%1 = %0
```

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

Author: krzysz00

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)
   }];

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);
 }
 
 //===----------------------------------------------------------------------===//

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
+}