[mlir][vector] Add canonicalization patterns for extractMap/insertMap

[ThomasRaoux]

Summary

Add basic canonicalization patterns for the extractMap/insertMap to allow them to be folded into Transfer ops.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2471

It may be useful to add the map attribute to the ops and already here fail if the map is not 1-D minor identity or something.
Problem is that the map is not yet designed so we are not 100% sure what form it will take.

Do what you think is best :)

2472

If you already pulled in EDSC for operator+ you might as well: vector_transfer_read(extract.getType() ...).

2494

same comment re vector_transfer_write

mlir/test/Dialect/Vector/vector-distribution.mlir

24

vector_add_read_write_at_id ?

Just trying to convey that this will only change 1 element in the vector.

mlir/test/lib/Transforms/TestVectorTransforms.cpp

132

For affine_apply ?

Yet more evidence, affine_apply is a math dialect concept and not an "affine dialect" concept.

132

s/math/arith

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2502

just curious why you did not make this a regular folding optimization?
(which kicks in at many more places and thus does not rely on populating this rewriting pattern at specific places)

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2471

True, I'll wait to insert the map if it's okay. I'll introduce it in it's own patch.

2472

Changed it.

2502

Good point, I moved it to ExtractMap folder.

mlir/test/lib/Transforms/TestVectorTransforms.cpp

132

Yes..

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2502

(I think) Aart's point is that it could even be part of the op's folding mechanisms ?
In practice it would be a canonicalizer because there is a type change.
So it would go to vector/Ops.cpp and be accessible via getCanonicalizationPatterns.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2502

I believe Aart was talking about the case where ExtractMap's source comes from an insertMap and they cancel each other. I moved it to extractMap folder. This just removes both instruction so it seems like it belongs to the folder, the type of the source of insert has to map the return of extract by construction (since the multiplicty matches)

%b = insertMap %a, %id, 32
extractMap %b, %id, 32
->
%a

I'll push it like that. I can change it post-commit if you think this is needed.

[nicolasvasilache]

SGTM :)

[aartbik]

Yes, I am happy with this as folder, like you changed now. That looks much cleaner to me.

[ThomasRaoux]

Revert marking TransferReadOp as readmem as some unit tests are not ready and have dead transfer_read

mlir/include/mlir/Dialect/Vector/VectorOps.td

def Vector_ExtractStridedSliceOp :
   let assemblyFormat = "$vector attr-dict `:` type($vector) `to` type(results)";
 }
 
 def Vector_TransferReadOp :
   Vector_Op<"transfer_read", [
       DeclareOpInterfaceMethods<VectorTransferOpInterface>,
       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
     ]>,
-    Arguments<(ins AnyMemRef:$memref, Variadic<Index>:$indices,
+    Arguments<(ins Arg<AnyMemRef, "", [MemRead]>:$memref,
+               Variadic<Index>:$indices,
                AffineMapAttr:$permutation_map, AnyType:$padding,
                OptionalAttr<BoolArrayAttr>:$masked)>,
     Results<(outs AnyVector:$vector)> {
 
   let summary = "Reads a supervector from memory into an SSA vector value.";
 
   let description = [{
     The `vector.transfer_read` op performs a read from a slice within a

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 //===- VectorTransforms.cpp - Conversion within the Vector dialect --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements target-independent rewrites as 1->N patterns.
 //
 //===----------------------------------------------------------------------===//
 
 #include <type_traits>
 
+#include "mlir/Dialect/Affine/EDSC/Builders.h"
 #include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/SCF/EDSC/Intrinsics.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
mlir::vector::distributPointwiseVectorOp(OpBuilder &builder, Operation *op,
   DistributeOps ops;
   ops.extract =
       builder.create<vector::ExtractMapOp>(loc, result, id, multiplicity);
   ops.insert =
       builder.create<vector::InsertMapOp>(loc, ops.extract, id, multiplicity);
   return ops;
 }
 
+struct TransferReadExtractPattern
+    : public OpRewritePattern<vector::ExtractMapOp> {
+  TransferReadExtractPattern(MLIRContext *context)
+      : OpRewritePattern<vector::ExtractMapOp>(context) {}
+  LogicalResult matchAndRewrite(vector::ExtractMapOp extract,
+                                PatternRewriter &rewriter) const override {
+    auto read = extract.vector().getDefiningOp<vector::TransferReadOp>();
+    if (!read)
+      return failure();
+    edsc::ScopedContext scope(rewriter, extract.getLoc());
+    using mlir::edsc::op::operator+;
+    SmallVector<Value, 4> indices(read.indices().begin(), read.indices().end());
+    indices.back() = indices.back() + extract.id();
+    Value newRead = rewriter.create<vector::TransferReadOp>(
+        read.getLoc(), extract.getType(), read.memref(), indices,
+        read.permutation_map(), read.padding(), ArrayAttr());
+    rewriter.replaceOp(extract, newRead);
+    return success();
+  }
+};
+
+struct TransferWriteInsertPattern
+    : public OpRewritePattern<vector::TransferWriteOp> {
+  TransferWriteInsertPattern(MLIRContext *context)
+      : OpRewritePattern<vector::TransferWriteOp>(context) {}
+  LogicalResult matchAndRewrite(vector::TransferWriteOp write,
+                                PatternRewriter &rewriter) const override {
+    auto insert = write.vector().getDefiningOp<vector::InsertMapOp>();
+    if (!insert)
+      return failure();
+    edsc::ScopedContext scope(rewriter, write.getLoc());
+    using mlir::edsc::op::operator+;
+    SmallVector<Value, 4> indices(write.indices().begin(),
+                                  write.indices().end());
+    indices.back() = indices.back() + insert.id();
+    rewriter.create<vector::TransferWriteOp>(
+        write.getLoc(), insert.vector(), write.memref(), indices,
+        write.permutation_map(), ArrayAttr());
+    rewriter.eraseOp(write);
+    return success();
+  }
+};
+
+struct ExtractMapFolderOp : public OpRewritePattern<vector::ExtractMapOp> {
+  ExtractMapFolderOp(MLIRContext *context)
+      : OpRewritePattern<vector::ExtractMapOp>(context) {}
+  LogicalResult matchAndRewrite(vector::ExtractMapOp extract,
+                                PatternRewriter &rewriter) const override {
+    auto insert = extract.vector().getDefiningOp<vector::InsertMapOp>();
+    if (!insert)
+      return failure();
+    if (extract.multiplicity() != insert.multiplicity() ||
+        extract.id() != insert.id())
+      return failure();
+    rewriter.replaceOp(extract, insert.vector());
+    return success();
+  }
+};
+
 // TODO: Add pattern to rewrite ExtractSlices(ConstantMaskOp).
 // TODO: Add this as DRR pattern.
 void mlir::vector::populateVectorToVectorTransformationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   // clang-format off
   patterns.insert<ShapeCastOpDecomposer,
                   ShapeCastOpFolder,
                   SplitTransferReadOp,
                   SplitTransferWriteOp,
-                  TupleGetFolderOp>(context);
+                  TupleGetFolderOp,
+                  TransferReadExtractPattern,
+                  TransferWriteInsertPattern,
+                  ExtractMapFolderOp>(context);
   // clang-format on
 }
 
 void mlir::vector::populateVectorSlicesLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<ExtractSlicesOpLowering, InsertSlicesOpLowering>(context);
 }
 

mlir/test/Dialect/Vector/vector-distribution.mlir

 // RUN: mlir-opt %s -test-vector-distribute-patterns | FileCheck %s
 
 // CHECK-LABEL: func @distribute_vector_add
 //  CHECK-SAME: (%[[ID:.*]]: index
 //  CHECK-NEXT:    %[[EXA:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
 //  CHECK-NEXT:    %[[EXB:.*]] = vector.extract_map %{{.*}}[%[[ID]] : 32] : vector<32xf32> to vector<1xf32>
 //  CHECK-NEXT:    %[[ADD:.*]] = addf %[[EXA]], %[[EXB]] : vector<1xf32>
 //  CHECK-NEXT:    %[[INS:.*]] = vector.insert_map %[[ADD]], %[[ID]], 32 : vector<1xf32> to vector<32xf32>
 //  CHECK-NEXT:    return %[[INS]] : vector<32xf32>
 func @distribute_vector_add(%id : index, %A: vector<32xf32>, %B: vector<32xf32>) -> vector<32xf32> {
   %0 = addf %A, %B : vector<32xf32>
   return %0: vector<32xf32>
 }
+
+// CHECK-LABEL: func @vector_add_read_write
+//  CHECK-SAME: (%[[ID:.*]]: index
+//       CHECK:    %[[EXA:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, vector<1xf32>
+//  CHECK-NEXT:    %[[EXB:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, vector<1xf32>
+//  CHECK-NEXT:    %[[ADD1:.*]] = addf %[[EXA]], %[[EXB]] : vector<1xf32>
+//  CHECK-NEXT:    %[[EXC:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %{{.*}} : memref<32xf32>, vector<1xf32>
+//  CHECK-NEXT:    %[[ADD2:.*]] = addf %[[ADD1]], %[[EXC]] : vector<1xf32>
+//  CHECK-NEXT:    vector.transfer_write %[[ADD2]], %{{.*}}[%{{.*}}] : vector<1xf32>, memref<32xf32>
+//  CHECK-NEXT:    return
+func @vector_add_read_write(%id : index, %A: memref<32xf32>, %B: memref<32xf32>, %C: memref<32xf32>, %D: memref<32xf32>) {
+  %c0 = constant 0 : index
+  %cf0 = constant 0.0 : f32
+  %a = vector.transfer_read %A[%c0], %cf0: memref<32xf32>, vector<32xf32>
+  %b = vector.transfer_read %B[%c0], %cf0: memref<32xf32>, vector<32xf32>
+  %acc = addf %a, %b: vector<32xf32>
+  %c = vector.transfer_read %C[%c0], %cf0: memref<32xf32>, vector<32xf32>
+  %d = addf %acc, %c: vector<32xf32>
+  vector.transfer_write %d, %D[%c0]: vector<32xf32>, memref<32xf32>
+  return
+}

mlir/test/lib/Transforms/TestVectorTransforms.cpp

void runOnFunction() override {
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorDistributePatterns
     : public PassWrapper<TestVectorDistributePatterns, FunctionPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<VectorDialect>();
+    registry.insert<AffineDialect>();
   }
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
     FuncOp func = getFunction();
     func.walk([&](AddFOp op) {
       OpBuilder builder(op);
       Optional<mlir::vector::DistributeOps> ops = distributPointwiseVectorOp(
           builder, op.getOperation(), func.getArgument(0), 32);
       assert(ops.hasValue());
       SmallPtrSet<Operation *, 1> extractOp({ops->extract});
       op.getResult().replaceAllUsesExcept(ops->insert.getResult(), extractOp);
     });
     patterns.insert<PointwiseExtractPattern>(ctx);
+    populateVectorToVectorTransformationPatterns(patterns, ctx);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorTransferFullPartialSplitPatterns
     : public PassWrapper<TestVectorTransferFullPartialSplitPatterns,
                          FunctionPass> {
   TestVectorTransferFullPartialSplitPatterns() = default;