[Mlir-commits] [mlir] 0c1ff26 - [mlir] [affine] add canonicalization for affine.vector_load, vector_store

[Uday Bondhugula]

Author: eopXD
Date: 2021-05-02T09:06:46+05:30
New Revision: 0c1ff26bd3df2e8732b5728b663efd7ec6f31a02

URL: https://github.com/llvm/llvm-project/commit/0c1ff26bd3df2e8732b5728b663efd7ec6f31a02
DIFF: https://github.com/llvm/llvm-project/commit/0c1ff26bd3df2e8732b5728b663efd7ec6f31a02.diff

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

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

Reviewed By: bondhugula

Differential Revision: https://reviews.llvm.org/D101691

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index e15b81d2676af..b1282ec8a1257 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -967,6 +967,8 @@ def AffineVectorLoadOp : AffineLoadOpBase<"vector_load"> {
       return result().getType().cast<VectorType>();
     }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
@@ -1029,6 +1031,8 @@ def AffineVectorStoreOp : AffineStoreOpBase<"vector_store"> {
       return value().getType().cast<VectorType>();
     }
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 #endif // AFFINE_OPS

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f295b0fd6eb65..62e1a66353458 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -853,15 +853,18 @@ struct SimplifyAffineOp : public OpRewritePattern<AffineOpTy> {
 
   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();
@@ -895,6 +898,22 @@ void SimplifyAffineOp<AffineStoreOp>::replaceAffineOp(
   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>
@@ -3267,6 +3286,11 @@ void AffineVectorLoadOp::build(OpBuilder &builder, OperationState &result,
   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();
@@ -3353,6 +3377,10 @@ void AffineVectorStoreOp::build(OpBuilder &builder, OperationState &result,
       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) {

diff  --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 7f528fb26fe5a..679e1225e3027 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -901,3 +901,27 @@ func @compose_affine_for_bounds(%N: index) {
   }
   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
+}