[mlir] Add primitive transform pattern to rewrite linalg.copy into vector.broadcast form

[tetuante]

Summary

This diff adds a transformation patter to rewrite linalg.copy as
broadcasting a vector into a vector.
It uses the same preconditioning as matmul (memory is contiguous).

[jdoerfert]

There should be a *-commits list in the subscribers.

[asaadaldien]

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

215

nit: s/fill/copy

[nicolasvasilache]

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

215

typo: fill -> copy

219

I think there are a few issues with this scheme:

1. this ignores the permutations on the copy so it will generate wrong code
2. linalg.copy itself is meant to copy between a memref with static behavior and a memref with "less static" behavior. In practice we should have logic for masking etc and there is some trickiness involved.
3. auto resVec = vector_broadcast(dstVec, srcVec); does not seem right

I would recommend instead evolving this revisions into a rewrite of linalg.copy into the proper vector.transfer_read and vector.transfer_write with proper permutations.
Note that the condition here is that only the source or target is fully static (not both, unlike linalg.fill and matmul).
Then in a separate revision we can start implementing rewrite and canonicalization patterns for transfers + splat and other things into the proper load / store + masks.

We may have to introduce first order masked loads too and delagate sone of the logic to LLVM, not too sure at this point.

[tetuante]

Abandoning this revision to evolve it into a rewrite of linalg.copy into the proper vector.transfer_read and vector.transfer_write with proper permutations.

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

LogicalResult mlir::linalg::vectorizeLinalgOpPrecondition(Operation *op) {
   auto linalgOp = cast<linalg::LinalgOp>(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::MatmulOp>(op) || isa<linalg::FillOp>(op))
+  if (isa<linalg::MatmulOp>(op) || isa<linalg::FillOp>(op) ||
+      isa<linalg::CopyOp>(op))
     return success();
 
   auto genericOp = dyn_cast<linalg::GenericOp>(op);
   if (!genericOp || !isMatmul(genericOp))
     return failure();
 
   // TODO(ntv): non-identity layout.
   auto isStaticMemRefWithIdentityLayout = [](Value v) {
if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
     // Vectorize fill as a vector.broadcast.
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
                          "]: Rewrite linalg.fill as vector.broadcast: "
                       << *op << ":\n");
     auto dstMemrefVec = vector_type_cast(fillOp.getOutputBuffer(0));
     auto dstVec = std_load(dstMemrefVec);
     auto resVec = vector_broadcast(dstVec, fillOp.value());
     std_store(resVec, dstMemrefVec);
+  } else if (auto copyOp = dyn_cast<linalg::CopyOp>(op)) {
+    // Vectorize fill as a vector.broadcast.
+    LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
+                         "]: Rewrite linalg.copy as vector.broadcast: "
+                      << *op << ":\n");
+    auto dstMemrefVec = vector_type_cast(copyOp.getOutputBuffer(0));
+    auto dstVec = std_load(dstMemrefVec);
+    auto srcMemrefVec = vector_type_cast(copyOp.getInput(0));
+    auto srcVec = std_load(srcMemrefVec);
+    auto resVec = vector_broadcast(dstVec, srcVec);
+    std_store(resVec, dstMemrefVec);
   } else {
     // Vectorize other ops as vector contraction (currently only matmul).
     LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
                          "]: Rewrite linalg op as vector.contract: "
                       << *op << ":\n");
     auto vA = std_load(vector_type_cast(linalgOp.getInput(0)));
     auto vB = std_load(vector_type_cast(linalgOp.getInput(1)));
     auto vectorMemRefC = vector_type_cast(linalgOp.getOutputBuffer(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: vector.broadcast {{.*}} : vector<8x16xf32> to vector<8x16xf32>
+
 func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
           %d = mulf %a, %b: f32
           %e = addf %c, %d: f32
           return %e: f32
         }
 #matmul_accesses = [
           affine_map<(m, n, k) -> (m, k)>,
           affine_map<(m, n, k) -> (k, n)>,

mlir/test/lib/DeclarativeTransforms/TestLinalgTransformPatterns.td

def : Pattern<(MatmulOp:$op $_, $_, $_),
                 PreconditionVectorizeLinalgOp
                ]>>)]>;
 def : Pattern<(FillOp:$op $_, $_),
               [(VectorizeLinalgOp)],
               [(Constraint<And<[
                 HasLinalgTransformMarker<"VECTORIZE">,
                 PreconditionVectorizeLinalgOp
                ]>>)]>;
+def : Pattern<(CopyOp:$op $_, $_, $_, $_),
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
+               ]>>)]>;
 def : Pattern<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
               [(VectorizeLinalgOp)],
               [(Constraint<And<[
                 HasLinalgTransformMarker<"VECTORIZE">,
                 PreconditionVectorizeLinalgOp
                ]>>)]>;