[mlir][Linalg] Generalize Vectorization of Linalg contractions

[nicolasvasilache]

Summary

This revision adds support for vectorizing named and generic contraction ops to vector.contract. Cases in which the memref is 0-D are special cased to emit std.load/std.store instead of vector.transfer. Relevant tests are added.

[aartbik]

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

96

can't we just

return success(isContraction(op))

at this point?

[aartbik]

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

3

Having several RUN command (and possible prefixes) makes sense if you want to run the *same* test and check it with different lowering strategies. However, here you seem to simply append new tests.

Wouldn't it make more sense to just start a new test file for this, to avoid running commands without any CHECKs. Or is the duplicate run intentional (if so, probably document that)?

[ftynse]

mlir/include/mlir/IR/Attributes.h

289

UnderlyingTy looks unused.

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

74

return success(/*condition here*/); ?

[nicolasvasilache]

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

3

I don't think we have particular best practices here.
I personally find it useful to "have more tests" and see that things run properly in even outside of the groomed positive tests.
But I can also see how this could be viewed as overhead and I am open to splitting this.
If we do the split, we should do it more systematically.
One could argue that the LLVM model should be the way (i.e. never have any test with split-input-files) but I would not be enthused by such a change.

[nicolasvasilache]

Format

[nicolasvasilache]

Rebase

mlir/include/mlir/IR/Attributes.h

private:
   };
 
 public:
   template <typename AttrTy>
   llvm::iterator_range<attr_value_iterator<AttrTy>> getAsRange() {
     return llvm::make_range(attr_value_iterator<AttrTy>(begin()),
                             attr_value_iterator<AttrTy>(end()));
   }
+  template <typename AttrTy, typename UnderlyingTy>
+  auto getAsRange() {
+    return llvm::map_range(getAsRange<AttrTy>(),
+                           [](AttrTy attr) { return attr.getValue(); });
+  }
 };
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// NamedAttribute is used for dictionary attributes, it holds an identifier for
 /// the name and a value for the attribute. The attribute pointer should always

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

 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 using namespace mlir::linalg;
 
 using llvm::dbgs;
 
 #define DEBUG_TYPE "linalg-vectorization"
 
-static bool hasMultiplyAddBody(linalg::GenericOp op) {
-  auto &r = op.region();
+static bool hasMultiplyAddBody(Region &r) {
   if (!llvm::hasSingleElement(r))
     return false;
   if (!llvm::hasNItems(r.front().begin(), r.front().end(), 3))
     return false;
 
   using mlir::matchers::m_Val;
   auto a = m_Val(r.front().getArgument(0));
   auto b = m_Val(r.front().getArgument(1));
   auto c = m_Val(r.front().getArgument(2));
   // TODO: Update this detection once we have  matcher support for specifying
   // that any permutation of operands matches.
   auto pattern1 = m_Op<YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(a, b), c));
   auto pattern2 = m_Op<YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(a, b)));
   auto pattern3 = m_Op<YieldOp>(m_Op<AddFOp>(m_Op<MulFOp>(b, a), c));
   auto pattern4 = m_Op<YieldOp>(m_Op<AddFOp>(c, m_Op<MulFOp>(b, a)));
   return pattern1.match(&r.front().back()) ||
          pattern2.match(&r.front().back()) ||
          pattern3.match(&r.front().back()) || pattern4.match(&r.front().back());
 }
 
 // TODO: Should be Tablegen'd from a single source that generates the op itself.
-static bool isRowMajorMatmul(linalg::GenericOp genericOp) {
-  return genericOp.getNumInputs() == 2 && genericOp.getNumOutputs() == 1 &&
-         isRowMajorMatmul(genericOp.indexing_maps()) &&
-         hasMultiplyAddBody(genericOp);
+static LogicalResult isContraction(Operation *op) {
+  // TODO: interface for named ops.
+  if (isa<linalg::BatchMatmulOp, linalg::MatmulOp, linalg::MatvecOp,
+          linalg::DotOp>(op))
+    return success();
+
+  auto genericOp = dyn_cast<linalg::GenericOp>(op);
+  if (!genericOp)
+    return failure();
+
+  auto mapRange =
+      genericOp.indexing_maps().getAsRange<AffineMapAttr, AffineMap>();
+
+  return (genericOp.getNumInputs() == 2 && genericOp.getNumOutputs() == 1 &&
+          llvm::all_of(
+              mapRange,
+              [](AffineMap m) { return m.isProjectedPermutation(); }) &&
+          hasMultiplyAddBody(genericOp.region()))
+             ? success()
+             : failure();
 }
 
-// TODO: 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, linalg::FillOp>(op))
+
+  if (isa<linalg::FillOp>(op))
     return success();
 
-  auto genericOp = dyn_cast<linalg::GenericOp>(op);
-  if (!genericOp || !::isRowMajorMatmul(genericOp))
-    return failure();
+  if (succeeded(isContraction(op)))
+    return success();
 
-  // 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;
-  };
-  return success(llvm::all_of(genericOp.getInputsAndOutputBuffers(),
-                              isStaticMemRefWithIdentityLayout));
+  return failure();
 }
 
 void mlir::linalg::vectorizeLinalgOp(OpBuilder &builder, Operation *op) {
   assert(succeeded(vectorizeLinalgOpPrecondition(op)));
 
-  if (auto convOp = dyn_cast<linalg::ConvOp>(op)) {
-    // TODO: add a level of indirection to linalg.generic.
-    if (convOp.padding())
-      llvm_unreachable("Unexpected conv with padding");
-  }
-
   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;
   }
 
-  // Vectorize other ops as vector contraction (currently only matmul).
+  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 VectorType::get(mt.getShape(), mt.getElementType());
+    return mt.getShape().empty()
+               ? VectorType()
+               : VectorType::get(mt.getShape(), mt.getElementType());
   };
   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);
-  SmallVector<Value, 4> indicesA(linalgOp.getInputShapedType(0).getRank(),
-                                 zero);
-  SmallVector<Value, 4> indicesB(linalgOp.getInputShapedType(1).getRank(),
-                                 zero);
-  SmallVector<Value, 4> indicesC(linalgOp.getOutputShapedType(0).getRank(),
-                                 zero);
-  Value a = vector_transfer_read(extractVectorTypeFromScalarView(viewA), viewA,
-                                 indicesA);
-  Value b = vector_transfer_read(extractVectorTypeFromScalarView(viewB), viewB,
-                                 indicesB);
-  Value c = vector_transfer_read(extractVectorTypeFromScalarView(viewC), viewC,
-                                 indicesC);
+  SmallVector<Value, 4> indicesA, indicesB, indicesC;
+  if (vtA)
+    indicesA = SmallVector<Value, 4>(vtA.getRank(), zero);
+  if (vtB)
+    indicesB = SmallVector<Value, 4>(vtB.getRank(), zero);
+  if (vtC)
+    indicesC = SmallVector<Value, 4>(vtC.getRank(), zero);
+  Value a = vtA ? vector_transfer_read(vtA, viewA, indicesA).value
+                : std_load(viewA, indicesA).value;
+  Value b = vtB ? vector_transfer_read(vtB, viewB, indicesB).value
+                : std_load(viewB, indicesB).value;
+  Value c = vtC ? vector_transfer_read(vtC, viewC, indicesC).value
+                : std_load(viewC, indicesC).value;
   Value res = vector_contract(a, b, c, linalgOp.indexing_maps(),
                               linalgOp.iterator_types());
+  if (vtC)
     vector_transfer_write(res, viewC, indicesC);
+  else
+    std_store(res, viewC, indicesC);
 }
 
 /// Check whether there is any interleaved use of any `values` between `firstOp`
 /// and `secondOp`. Conservatively return `true` if any op or value is in a
 /// different block.
 static bool mayExistInterleavedUses(Operation *firstOp, Operation *secondOp,
                                     ValueRange values) {
   StringRef dbgPref = "\n[" DEBUG_TYPE "]: ";

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 %A, %B, %C {__internal_linalg_transform__ = "START"} :
     (memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
      memref<1584x1584xf32, offset: 0, strides: [1584, 1]>,
      memref<1584x1584xf32, offset: 0, strides: [1584, 1]>)

 //      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>
 //
 //      CHECK: linalg.copy
+
+// VECTOR-CONTRACTION-LABEL: contraction_dot
+func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
+  // VECTOR-CONTRACTION: vector.contract
+  // VECTOR-CONTRACTION-SAME: vector<1584xf32>, vector<1584xf32> into f32
+  linalg.dot(%A, %B, %C) : memref<1584xf32>, memref<1584xf32>, memref<f32>
+  return
+}
+
+// VECTOR-CONTRACTION-LABEL: contraction_matvec
+func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) {
+  // VECTOR-CONTRACTION: vector.contract
+  // VECTOR-CONTRACTION-SAME: vector<1584x1584xf32>, vector<1584xf32> into vector<1584xf32>
+  linalg.matvec %A, %B, %C :
+    (memref<1584x1584xf32>, memref<1584xf32>, memref<1584xf32>)
+  return
+}
+
+// VECTOR-CONTRACTION-LABEL: contraction_matmul
+func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
+  // VECTOR-CONTRACTION: vector.contract
+  // VECTOR-CONTRACTION-SAME: vector<1584x1584xf32>, vector<1584x1584xf32> into vector<1584x1584xf32>
+  linalg.matmul %A, %B, %C :
+    (memref<1584x1584xf32>, memref<1584x1584xf32>, memref<1584x1584xf32>)
+  return
+}
+
+// VECTOR-CONTRACTION-LABEL: contraction_batch_matmul
+func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) {
+  // VECTOR-CONTRACTION: vector.contract
+  // VECTOR-CONTRACTION-SAME: vector<1584x1584x1584xf32>, vector<1584x1584x1584xf32> into vector<1584x1584x1584xf32>
+  linalg.batch_matmul %A, %B, %C :
+    (memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>)
+  return
+}

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

struct TestLinalgTransforms
   Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
                                     llvm::cl::desc("Test promotion options"),
                                     llvm::cl::init(false)};
   Option<bool> testVectorTransferForwardingPatterns{
       *this, "test-vector-transfer-forwarding-patterns",
       llvm::cl::desc(
           "Test a fused pass that forwards linalg.copy to vector.transfer"),
       llvm::cl::init(false)};
+  Option<bool> testGenericToVectorPattern{
+      *this, "test-contraction-to-vector-patterns",
+      llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
+                     "in vector.contract form"),
+      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//

 
 static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
   OwningRewritePatternList forwardPattern;
   forwardPattern.insert<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
   forwardPattern.insert<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
   applyPatternsAndFoldGreedily(funcOp, forwardPattern);
 }
 
+static void applyContractionToVectorPatterns(FuncOp funcOp) {
+  OwningRewritePatternList patterns;
+  patterns.insert<LinalgVectorizationPattern<BatchMatmulOp>,
+                  LinalgVectorizationPattern<MatmulOp>,
+                  LinalgVectorizationPattern<MatvecOp>,
+                  LinalgVectorizationPattern<DotOp>,
+                  LinalgVectorizationPattern<GenericOp>>(funcOp.getContext());
+  applyPatternsAndFoldGreedily(funcOp, patterns);
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
   auto lambda = [&](void *) {
     getFunction().walk([](LinalgOp op) {
       op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
   };
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
 
   if (testPromotionOptions) {
     OwningRewritePatternList patterns;
     fillPromotionCallBackPatterns(&getContext(), patterns);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
     return;
   }
   if (testPatterns)
     return applyPatterns(getFunction());
   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
     return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
+  if (testGenericToVectorPattern)
+    return applyContractionToVectorPatterns(getFunction());
 }
 
 namespace mlir {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace mlir