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

[asaadaldien]

Summary

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

[rriddle]

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

214

nit: if (auto fillOp = dyn_cast<linalg::FillOp(op)) {

[merge_guards_bot]

Unit tests: pass. 62193 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

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

216

fillas -> fill as

218

Period at end of comment.

[asaadaldien]

Resolve comments & clang-format

[nicolasvasilache]

Looks good modulo Aart's comments.

Additional request for next pattern: 3rd time's a charm, let's please start refactoring the rewrites and template them on op types.

[merge_guards_bot]

Unit tests: pass. 62193 tests passed, 0 failed and 815 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/LinalgTransforms.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/VectorOps/VectorOps.h"
 #include "mlir/EDSC/Helpers.h"
+#include "mlir/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
 
 #define DEBUG_TYPE "linalg-transforms"
 
 using namespace mlir;
 using namespace mlir::edsc;
static bool isMatmul(linalg::GenericOp genericOp) {
   auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}));
   auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}));
   auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}));
   auto maps = ArrayAttr::get({mapA, mapB, mapC}, ctx);
   return genericOp.getNumInputs() == 2 && genericOp.getNumOutputs() == 1 &&
          genericOp.indexing_maps() == maps && hasMultiplyAddBody(genericOp);
 }
 
-// TODO(ntv): This is in fact much more general than just vectorization for
-// matmul ops.
+// TODO(ntv, ataei): This is in fact much more general than just vectorization
+// for matmul and fill ops.
 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))
+  if (isa<linalg::MatmulOp>(op) || isa<linalg::FillOp>(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) {
     auto m = v.getType().dyn_cast<MemRefType>();
     if (!m || !m.hasStaticShape() || !m.getAffineMaps().empty())
       return false;
     return true;
   };
   if (!llvm::all_of(genericOp.getInputsAndOutputBuffers(),
                     isStaticMemRefWithIdentityLayout))
     return failure();
   return success();
 }
 
 SmallVector<Value, 0> mlir::linalg::vectorizeLinalgOp(PatternRewriter &rewriter,
                                                       Operation *op) {
-  LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE
-                       "]: Rewrite linalg op as vector.contract: "
-                    << *op << ":\n");
+  using edsc::intrinsics::std_load;
+  using edsc::intrinsics::std_store;
+  using vector_contract = edsc::intrinsics::ValueBuilder<vector::ContractionOp>;
+  using vector_broadcast = edsc::intrinsics::ValueBuilder<vector::BroadcastOp>;
+  using vector_type_cast = edsc::intrinsics::ValueBuilder<vector::TypeCastOp>;
 
   assert(succeeded(vectorizeLinalgOpPrecondition(op)) &&
          "DRR failure case must be a precondition");
-
   auto linalgOp = cast<linalg::LinalgOp>(op);
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   edsc::ScopedContext scope(rewriter, op->getLoc());
-  using edsc::intrinsics::std_load;
-  using edsc::intrinsics::std_store;
-  using vector_contract = edsc::intrinsics::ValueBuilder<vector::ContractionOp>;
-  using vector_type_cast = edsc::intrinsics::ValueBuilder<vector::TypeCastOp>;
+
+  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 {
+    // 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));
     auto vC = std_load(vectorMemRefC);
     auto vRes = vector_contract(vA, vB, vC, linalgOp.indexing_maps(),
                                 linalgOp.iterator_types());
     std_store(vRes, vectorMemRefC);
+  }
   return {};
 }
 
 //============================================================================//
 // Precondition and transformation for permutation of Linalg generic ops.
 //============================================================================//
 LogicalResult mlir::linalg::permuteGenericLinalgOpPrecondition(
     Operation *op, ArrayRef<unsigned> permutation) {

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

func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
   linalg.matmul(%A, %B, %C) { __internal_linalg_transform__ = "VECTORIZE"} :
     memref<8x16xf32>, memref<16x32xf32>, memref<8x32xf32>
   return
 }
 // CHECK-LABEL: func @vectorization_test_2
 //       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>
+
 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

 // Linalg to vector contraction patterns.
 //===----------------------------------------------------------------------===//
 def : Pattern<(MatmulOp:$op $_, $_, $_),
               [(VectorizeLinalgOp)],
               [(Constraint<And<[
                 HasLinalgTransformMarker<"VECTORIZE">,
                 PreconditionVectorizeLinalgOp
                ]>>)]>;
+def : Pattern<(FillOp:$op $_, $_),
+              [(VectorizeLinalgOp)],
+              [(Constraint<And<[
+                HasLinalgTransformMarker<"VECTORIZE">,
+                PreconditionVectorizeLinalgOp
+               ]>>)]>;
 def : Pattern<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
               [(VectorizeLinalgOp)],
               [(Constraint<And<[
                 HasLinalgTransformMarker<"VECTORIZE">,
                 PreconditionVectorizeLinalgOp
                ]>>)]>;
 
+
 //===----------------------------------------------------------------------===//
 // Linalg generic permutation patterns.
 //===----------------------------------------------------------------------===//
 def : Pat<(GenericOp:$op $_, $_, $_, $_, $_, $_, $_, $_),
               (PermuteGenericLinalgOp<[1, 2, 0], "PERMUTE"> $op),
               [(Constraint<And<[
                 HasNoLinalgTransformMarker,
                 AffineMapDomainHasDim<3>,