[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

2468

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 :)

2469

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

2491

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

2499

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

2468

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

2469

Changed it.

2499

Good point, I moved it to ExtractMap folder.

mlir/test/lib/Transforms/TestVectorTransforms.cpp

132

Yes..

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2499

(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

2499

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

let extraClassDeclaration = [{
     VectorType getResultType() {
       return getResult().getType().cast<VectorType>();
     }
   }];
   let assemblyFormat = [{
     $vector `[` $id `:` $multiplicity `]` attr-dict `:` type($vector) `to`
     type(results)
   }];
+
+  let hasFolder = 1;
 }
 
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [NoSideEffect,
                              AllTypesMatch<["lhs", "rhs", "acc", "result"]>]>,
     Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyVector:$acc)>,
     Results<(outs AnyVector:$result)> {
   let summary = "vector fused multiply-add";

mlir/lib/Dialect/Vector/VectorOps.cpp

if (op.getSourceVectorType().getShape().size() != 1 ||
     return op.emitOpError("expects source and destination vectors of rank 1");
   if (op.getResultType().getNumElements() * (int64_t)op.multiplicity() !=
       op.getSourceVectorType().getNumElements())
     return op.emitOpError("vector sizes mismatch. Source size must be equal "
                           "to destination size * multiplicity");
   return success();
 }
 
+OpFoldResult ExtractMapOp::fold(ArrayRef<Attribute> operands) {
+  auto insert = vector().getDefiningOp<vector::InsertMapOp>();
+  if (insert == nullptr || multiplicity() != insert.multiplicity() ||
+      id() != insert.id())
+    return {};
+  return insert.vector();
+}
+
 //===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(BroadcastOp op) {
   VectorType srcVectorType = op.getSourceType().dyn_cast<VectorType>();
   VectorType dstVectorType = op.getVectorType();
   // Scalar to vector broadcast is always valid. A vector

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::TransferReadOp> {
+  TransferReadExtractPattern(MLIRContext *context)
+      : OpRewritePattern<vector::TransferReadOp>(context) {}
+  LogicalResult matchAndRewrite(vector::TransferReadOp read,
+                                PatternRewriter &rewriter) const override {
+    if (!read.getResult().hasOneUse())
+      return failure();
+    auto extract =
+        dyn_cast<vector::ExtractMapOp>(*read.getResult().getUsers().begin());
+    if (!extract)
+      return failure();
+    edsc::ScopedContext scope(rewriter, read.getLoc());
+    using mlir::edsc::op::operator+;
+    using namespace mlir::edsc::intrinsics;
+    SmallVector<Value, 4> indices(read.indices().begin(), read.indices().end());
+    indices.back() = indices.back() + extract.id();
+    Value newRead = vector_transfer_read(extract.getType(), read.memref(),
+                                         indices, read.permutation_map(),
+                                         read.padding(), ArrayAttr());
+    newRead = rewriter.create<vector::InsertMapOp>(
+        read.getLoc(), newRead, extract.id(), extract.multiplicity());
+    rewriter.replaceOp(read, 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+;
+    using namespace mlir::edsc::intrinsics;
+    SmallVector<Value, 4> indices(write.indices().begin(),
+                                  write.indices().end());
+    indices.back() = indices.back() + insert.id();
+    vector_transfer_write(insert.vector(), write.memref(), indices,
+                          write.permutation_map(), ArrayAttr());
+    rewriter.eraseOp(write);
+    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>(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;