[mlir][vector] Make linalg FillOp vectorization use Transfer op

[mravishankar]

Just one question based on some questions.

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

122

Its not directly related to this change, but doesnt extractVectorTypeFromScalarView also need to check for static shapes?

[ThomasRaoux]

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

122

There is already a check in vectorizeLinalgOpPrecondition that all the shapes are static. So the pattern would have failed before.

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

 }
 
 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;
-  }
-
   // 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 fillOp = dyn_cast<linalg::FillOp>(op)) {
+    // Vectorize fill as a vector.broadcast.
+    LLVM_DEBUG(dbgs() << dbgPref
+                      << "Rewrite linalg.fill as vector.broadcast: " << *op);
+    Value viewOutput = fillOp.output();
+    if (VectorType outputType = extractVectorTypeFromScalarView(viewOutput)) {
+      auto vecType =
+          VectorType::get(fillOp.getOutputBufferType(0).getShape(),
+                          fillOp.getOutputBufferType(0).getElementType());
+      Value vector = vector_broadcast(vecType, fillOp.value());
+      Value zero = std_constant_index(0);
+      SmallVector<Value, 4> indicesOutput(outputType.getRank(), zero);
+      vector_transfer_write(vector, viewOutput, indicesOutput);
+    } else {
+      std_store(fillOp.value(), viewOutput);
+    }
+    return;
+  }
   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();

mlir/test/Dialect/Linalg/transform-patterns-matmul-to-vector.mlir

 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-1d | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-matmul-to-vector-patterns-tile-2d | FileCheck %s
 // RUN: mlir-opt %s -test-linalg-transform-patterns=test-contraction-to-vector-patterns | FileCheck %s --check-prefix=VECTOR-CONTRACTION
 
 func @matmul(%A: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                   %C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>) {
   linalg.matmul {__internal_linalg_transform__ = "START"}
     ins(%A, %B: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
                 memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)
    outs(%C: memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)
   return
 }
 
 // CHECK-LABEL:func @matmul
-//      CHECK: store {{.*}}[] : memref<vector<8x16xf32>>
-//      CHECK: store {{.*}}[] : memref<vector<16x12xf32>>
-//      CHECK: store {{.*}}[] : memref<vector<8x12xf32>>
+//      CHECK: vector.transfer_write {{.*}} : vector<8x16xf32>, memref<8x16xf32>
+//      CHECK: vector.transfer_write {{.*}} : vector<16x12xf32>, memref<16x12xf32>
+//      CHECK: vector.transfer_write {{.*}} : vector<8x12xf32>, memref<8x12xf32>
 //
 //      CHECK: linalg.copy
 //      CHECK: linalg.copy
 //      CHECK: linalg.copy
 //
 //      CHECK: vector.contract
 // CHECK-SAME:   iterator_types = ["parallel", "parallel", "reduction"]
 // CHECK-SAME:   : vector<8x16xf32>, vector<16x12xf32> into vector<8x12xf32>

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

 //       CHECK: vector.contract {{.*}} :
 //                vector<8x16xf32>, vector<16x32xf32> into vector<8x32xf32>
 
 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>
+//       CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32>
+//       CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
+
+func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) {
+  linalg.fill(%A, %arg0) { __internal_linalg_transform__ = "VECTORIZE"} :  memref<f32>, f32
+  return
+}
+// CHECK-LABEL: func @test_vectorize_fill
+//  CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[V:.*]]: f32)
+//       CHECK:   store %[[V]], %[[M]][] : memref<f32>
 
 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>