[mlir] [affine] add canonicalization for affine.vector_load, vector_store

[eopXD]

Summary

Added canonicalization for vector_load and vector_store. An existing
pattern SimplifyAffineOp can be reused to compose maps that supplies
result into them. Added AffineVectorStoreOp and AffineVectorLoadOp
into static_assert of SimplifyAffineOp to allow operation to use it.

This fixes the bug filed: https://bugs.llvm.org/show_bug.cgi?id=50058

[eopXD]

I've also added canonicalizeMapAndOperands into SimplifyAffineOp, this will let maps be canonicalized, allowing the maps operations to be folded into the operand.

[bondhugula]

Thanks for adding this. Looks great to me. Please do address the clang-format warnings.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

907

You can run git clang-format HEAD~ to fix all these (assume HEAD is at this commit).

mlir/test/Dialect/Affine/canonicalize.mlir

927

New line at end of file.

[eopXD]

Fix format warnings.

[eopXD]

Formats fixed, thanks for reviewing.
I don't have the access to commit to llvm, so may you help me to do it?

[bondhugula]

In D101691#2731306, @eopXD wrote:

Formats fixed, thanks for reviewing.
I don't have the access to commit to llvm, so may you help me to do it?

Done. (I fixed a minor whitespace issue in canonicalize.mlir at line #920 before committing.)

[eopXD]

In D101691#2731695, @bondhugula wrote:

In D101691#2731306, @eopXD wrote:

Formats fixed, thanks for reviewing.
I don't have the access to commit to llvm, so may you help me to do it?

Done. (I fixed a minor whitespace issue in canonicalize.mlir at line #920 before committing.)

Thank you for your modification and commit.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

OpBuilder<(ins "VectorType":$resultType, "Value":$memref,
       "AffineMap":$map, "ValueRange":$mapOperands)>
   ];
 
   let extraClassDeclaration = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return result().getType().cast<VectorType>();
     }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
   let summary = "affine vector store operation";
   let description = [{
     The "affine.vector_store" is the vector counterpart of
     [affine.store](#affinestore-affinestoreop). It writes a
     [vector](../LangRef.md#vector-type), supplied as its first operand,
OpBuilder<(ins "Value":$valueToStore, "Value":$memref, "AffineMap":$map,
       "ValueRange":$mapOperands)>
   ];
 
   let extraClassDeclaration = extraClassDeclarationBase # [{
     VectorType getVectorType() {
       return value().getType().cast<VectorType>();
     }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 #endif // AFFINE_OPS

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> {
 
   /// Replace the affine op with another instance of it with the supplied
   /// map and mapOperands.
   void replaceAffineOp(PatternRewriter &rewriter, AffineOpTy affineOp,
                        AffineMap map, ArrayRef<Value> mapOperands) const;
 
   LogicalResult matchAndRewrite(AffineOpTy affineOp,
                                 PatternRewriter &rewriter) const override {
-    static_assert(llvm::is_one_of<AffineOpTy, AffineLoadOp, AffinePrefetchOp,
-                                  AffineStoreOp, AffineApplyOp, AffineMinOp,
-                                  AffineMaxOp>::value,
-                  "affine load/store/apply/prefetch/min/max op expected");
+    static_assert(
+        llvm::is_one_of<AffineOpTy, AffineLoadOp, AffinePrefetchOp,
+                        AffineStoreOp, AffineApplyOp, AffineMinOp, AffineMaxOp,
+                        AffineVectorStoreOp, AffineVectorLoadOp>::value,
+        "affine load/store/vectorstore/vectorload/apply/prefetch/min/max op "
+        "expected");
     auto map = affineOp.getAffineMap();
     AffineMap oldMap = map;
     auto oldOperands = affineOp.getMapOperands();
     SmallVector<Value, 8> resultOperands(oldOperands);
     composeAffineMapAndOperands(&map, &resultOperands);
+    canonicalizeMapAndOperands(&map, &resultOperands);
     if (map == oldMap && std::equal(oldOperands.begin(), oldOperands.end(),
                                     resultOperands.begin()))
       return failure();
 
     replaceAffineOp(rewriter, affineOp, map, resultOperands);
     return success();
   }
 };

 }
 template <>
 void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(
     PatternRewriter &rewriter, AffineStoreOp store, AffineMap map,
     ArrayRef<Value> mapOperands) const {
   rewriter.replaceOpWithNewOp<AffineStoreOp>(
       store, store.getValueToStore(), store.getMemRef(), map, mapOperands);
 }
+template <>
+void SimplifyAffineOp<AffineVectorLoadOp>::replaceAffineOp(
+    PatternRewriter &rewriter, AffineVectorLoadOp vectorload, AffineMap map,
+    ArrayRef<Value> mapOperands) const {
+  rewriter.replaceOpWithNewOp<AffineVectorLoadOp>(
+      vectorload, vectorload.getVectorType(), vectorload.getMemRef(), map,
+      mapOperands);
+}
+template <>
+void SimplifyAffineOp<AffineVectorStoreOp>::replaceAffineOp(
+    PatternRewriter &rewriter, AffineVectorStoreOp vectorstore, AffineMap map,
+    ArrayRef<Value> mapOperands) const {
+  rewriter.replaceOpWithNewOp<AffineVectorStoreOp>(
+      vectorstore, vectorstore.getValueToStore(), vectorstore.getMemRef(), map,
+      mapOperands);
+}
 
 // Generic version for ops that don't have extra operands.
 template <typename AffineOpTy>
 void SimplifyAffineOp<AffineOpTy>::replaceAffineOp(
     PatternRewriter &rewriter, AffineOpTy op, AffineMap map,
     ArrayRef<Value> mapOperands) const {
   rewriter.replaceOpWithNewOp<AffineOpTy>(op, map, mapOperands);
 }
void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
   int64_t rank = memrefType.getRank();
   // Create identity map for memrefs with at least one dimension or () -> ()
   // for zero-dimensional memrefs.
   auto map =
       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
   build(builder, result, resultType, memref, map, indices);
 }
 
+void AffineVectorLoadOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                     MLIRContext *context) {
+  results.add<SimplifyAffineOp<AffineVectorLoadOp>>(context);
+}
+
 static ParseResult parseAffineVectorLoadOp(OpAsmParser &parser,
                                            OperationState &result) {
   auto &builder = parser.getBuilder();
   auto indexTy = builder.getIndexType();
 
   MemRefType memrefType;
   VectorType resultType;
   OpAsmParser::OperandType memrefInfo;
void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
   auto memrefType = memref.getType().cast<MemRefType>();
   int64_t rank = memrefType.getRank();
   // Create identity map for memrefs with at least one dimension or () -> ()
   // for zero-dimensional memrefs.
   auto map =
       rank ? builder.getMultiDimIdentityMap(rank) : builder.getEmptyAffineMap();
   build(builder, result, valueToStore, memref, map, indices);
 }
+void AffineVectorStoreOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<SimplifyAffineOp<AffineVectorStoreOp>>(context);
+}
 
 static ParseResult parseAffineVectorStoreOp(OpAsmParser &parser,
                                             OperationState &result) {
   auto indexTy = parser.getBuilder().getIndexType();
 
   MemRefType memrefType;
   VectorType resultType;
   OpAsmParser::OperandType storeValueInfo;

mlir/test/Dialect/Affine/canonicalize.mlir

 func @compose_affine_for_bounds(%N: index) {
   %u = affine.apply affine_map<(d0) -> (d0 + 2)>(%N)
   %l = affine.apply affine_map<(d0) -> (d0 - 2)>(%N)
   affine.for %i = %l to %u {
     "foo"() : () -> ()
   }
   return
 }
+
+// -----
+
+// Compose maps into affine.vector_load / affine.vector_store
+
+// CHECK-LABEL: func @compose_into_affine_vector_load_vector_store
+// CHECK: affine.for %[[IV:.*]] = 0 to 1024
+// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]] + 1]
+// CHECK-NEXT: affine.vector_store %{{.*}}, %{{.*}}[%[[IV]] + 1]
+// CHECK-NEXT: affine.vector_load %{{.*}}[%[[IV]]]
+func @compose_into_affine_vector_load_vector_store(%A : memref<1024xf32>, %u : index) {
+  affine.for %i = 0 to 1024 {
+    // Make sure the unused operand (%u below) gets dropped as well.
+    %idx = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i, %u)
+    %0 = affine.vector_load %A[%idx] : memref<1024xf32>, vector<8xf32>
+    affine.vector_store %0, %A[%idx] : memref<1024xf32>, vector<8xf32>
+
+    // Map remains the same, but operand changes on composition.
+    %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
+    %1 = affine.vector_load %A[%copy] : memref<1024xf32>, vector<8xf32>
+    "prevent.dce"(%1) : (vector<8xf32>) -> ()
+  }
+  return
+}