[mlir][linalg] Add vectorization transform for CopyOp

[ThomasRaoux]

Summary

CopyOp get vectorized to vector.transfer_read followed by vector.transfer_write.

[aartbik]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

122

can you please add a test case for this also? (copy that falls into scalar code)
since technically that is a new path through the code now

[ThomasRaoux]

Handle scalar case and add a test for it.

[ThomasRaoux]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

122

Thanks Aart. Added a test for it and fixed the code to make it work.

[nicolasvasilache]

Thanks @ThomasRaoux

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
   // All types must be static shape to go to vector.
   for (Value operand : linalgOp.getInputsAndOutputBuffers())
     if (!operand.getType().cast<ShapedType>().hasStaticShape())
       return failure();
   for (Type outputTensorType : linalgOp.getOutputTensorTypes())
     if (!outputTensorType.cast<ShapedType>().hasStaticShape())
       return failure();
 
-  if (isa<linalg::FillOp>(op))
+  if (isa<linalg::FillOp, linalg::CopyOp>(op))
     return success();
 
   return isContraction(op);
 }
 
 void mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op) {
   assert(succeeded(vectorizeLinalgOpPrecondition(op)));
 
   StringRef dbgPref = "\n[" DEBUG_TYPE "]: ";
   (void)dbgPref;
   edsc::ScopedContext scope(builder, op->getLoc());
   if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
     // Vectorize fill as a vector.broadcast.
     LLVM_DEBUG(dbgs() << dbgPref
                       << "Rewrite linalg.fill as vector.broadcast: " << *op);
     Value memref = vector_type_cast(fillOp.getOutputBuffer(0));
     Value dst = std_load(memref);
     Value res = vector_broadcast(dst.getType(), fillOp.value());
     std_store(res, memref);
     return;
   }
 
-  assert(succeeded(isContraction(op)) && "Expected contraction");
-
-  // Vectorize other ops as vector contraction.
-  // TODO: interface.
-  LLVM_DEBUG(dbgs() << dbgPref
-                    << "Rewrite linalg op as vector.contract: " << *op);
   // In the case of 0-D memrefs, return null and special case to scalar load or
   // store later.
   auto extractVectorTypeFromScalarView = [](Value v) {
     MemRefType mt = v.getType().cast<MemRefType>();
     return mt.getShape().empty()
                ? VectorType()
                : VectorType::get(mt.getShape(), mt.getElementType());
   };
+
+  if (auto copyOp = dyn_cast<linalg::CopyOp>(op)) {
+    // Vectorize copy as a vector.transfer_read+vector.transfer_write.
+    LLVM_DEBUG(dbgs() << dbgPref
+                      << "Rewrite linalg.copy as vector.transfer_read + "
+                         "vector.transfer_write: "
+                      << *op);
+    Value zero = std_constant_index(0);
+    Value viewInput = copyOp.input();
+    Value viewOutput = copyOp.output();
+    Value vector;
+    if (VectorType inputType = extractVectorTypeFromScalarView(viewInput)) {
+      SmallVector<Value, 4> indicesInput(inputType.getRank(), zero);
+      if (copyOp.inputPermutation())
+        vector = vector_transfer_read(
+            extractVectorTypeFromScalarView(viewInput), viewInput, indicesInput,
+            copyOp.inputPermutation().getValue());
+      else
+        vector =
+            vector_transfer_read(extractVectorTypeFromScalarView(viewInput),
+                                 viewInput, indicesInput);
+    } else {
+      vector = std_load(viewInput).value;
+    }
+    if (VectorType outputType = extractVectorTypeFromScalarView(viewOutput)) {
+      SmallVector<Value, 4> indicesOutput(outputType.getRank(), zero);
+      if (copyOp.outputPermutation())
+        vector_transfer_write(vector, viewOutput, indicesOutput,
+                              copyOp.outputPermutation().getValue());
+      else
+        vector_transfer_write(vector, viewOutput, indicesOutput);
+    } else {
+      std_store(vector, viewOutput);
+    }
+    return;
+  }
+
+  assert(succeeded(isContraction(op)) && "Expected contraction");
+
+  // Vectorize other ops as vector contraction.
+  // TODO: interface.
+  LLVM_DEBUG(dbgs() << dbgPref
+                    << "Rewrite linalg op as vector.contract: " << *op);
   auto linalgOp = cast<linalg::LinalgOp>(op);
   Value viewA = linalgOp.getInput(0);
   Value viewB = linalgOp.getInput(1);
   Value viewC = linalgOp.getOutputBuffer(0);
   VectorType vtA = extractVectorTypeFromScalarView(viewA);
   VectorType vtB = extractVectorTypeFromScalarView(viewB);
   VectorType vtC = extractVectorTypeFromScalarView(viewC);
   Value zero = std_constant_index(0);

mlir/test/Dialect/Linalg/transform-patterns.mlir

 
 func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) {
   linalg.fill(%A, %arg0) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<8x16xf32>, f32
   return
 }
 // CHECK-LABEL: func @test_vectorize_fill
 //       CHECK: vector.broadcast {{.*}} : f32 to vector<8x16xf32>
 
+func @test_vectorize_copy(%A : memref<8x16xf32>, %B : memref<8x16xf32>) {
+  linalg.copy(%A, %B) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<8x16xf32>, memref<8x16xf32>
+  return
+}
+// CHECK-LABEL: func @test_vectorize_copy
+//       CHECK: %[[V:.*]] = vector.transfer_read {{.*}} : memref<8x16xf32>, vector<8x16xf32>
+//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
+
+func @test_vectorize_copy_scalar(%A : memref<f32>, %B : memref<f32>) {
+  linalg.copy(%A, %B) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<f32>, memref<f32>
+  return
+}
+// CHECK-LABEL: func @test_vectorize_copy_scalar
+//       CHECK: %[[V:.*]] = load {{.*}} : memref<f32>
+//       CHECK: store %[[V]], {{.*}} : memref<f32>
+
+
 #matmul_accesses = [
   affine_map<(m, n, k) -> (m, k)>,
   affine_map<(m, n, k) -> (k, n)>,
   affine_map<(m, n, k) -> (m, n)>
 ]
 #generic_matmul_trait = {
   args_in = 2,
   args_out = 1,

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

patterns.insert<LinalgLoweringPattern<DotOp>>(
       /*loweringType=*/LinalgLoweringType::Loops,
       LinalgMarker(Identifier::get("REG", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg to vector contraction patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgVectorizationPattern<MatmulOp>,
                   LinalgVectorizationPattern<FillOp>,
+                  LinalgVectorizationPattern<CopyOp>,
                   LinalgVectorizationPattern<GenericOp>>(
       ctx, LinalgMarker(Identifier::get("VECTORIZE", ctx)));
 
   //===--------------------------------------------------------------------===//
   // Linalg generic permutation patterns.
   //===--------------------------------------------------------------------===//
   patterns.insert<LinalgInterchangePattern<GenericOp>>(
       ctx,