[mlir] [VectorOps] Initial framework for progressively lowering vector.contract

[aartbik]

Summary

This sets the basic framework for lowering vector.contract progressively
into simpler vector.contract operations until a direct vector.reduction
operation is reached. More details will be filled out progressively as well.

[nicolasvasilache]

Minor nits, approving conditionally.
Thanks Aart!

mlir/lib/Dialect/VectorOps/VectorOps.cpp

364

trivial braces

373

trivial braces

mlir/test/lib/Transforms/TestVectorTransforms.cpp

56

lowerCamelCase names here and below plz.
It seems we missed a beat with slicesPass in the past.

[aartbik]

addressed comments, rebased

mlir/include/mlir/Dialect/VectorOps/VectorOps.h

 /// ops in terms of more elementary vector "slice" ops. If all
 /// "produced" tuple values are "consumed" (the most common
 /// use for "slices" ops), this lowering removes all tuple related
 /// operations as well (through DCE and folding). If tuple values
 /// "leak" coming in, however, some tuple related ops will remain.
 void populateVectorSlicesLoweringPatterns(OwningRewritePatternList &patterns,
                                           MLIRContext *context);
 
+/// Collect a set of vector contraction transformation patterns
+/// that express all vector.contract ops in terms of more elementary
+/// extraction and reduction ops.
+void populateVectorContractLoweringPatterns(OwningRewritePatternList &patterns,
+                                            MLIRContext *context);
+
 /// Returns the integer type required for subscripts in the vector dialect.
 IntegerType getVectorSubscriptType(Builder &builder);
 
 /// Returns an integer array attribute containing the given values using
 /// the integer type required for subscripts in the vector dialect.
 ArrayAttr getVectorSubscriptAttr(Builder &b, ArrayRef<int64_t> values);
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/VectorOps/VectorOps.h.inc"
 
 } // end namespace vector
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_VECTOROPS_VECTOROPS_H

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

let extraClassDeclaration = [{
     void getIterationIndexMap(
       std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap);
 
     std::vector<std::pair<int64_t, int64_t>> getContractingDimMap();
     std::vector<std::pair<int64_t, int64_t>> getBatchDimMap();
   }];
 }
 
+def Vector_ReductionOp :
+  Vector_Op<"reduction", [NoSideEffect,
+     PredOpTrait<"source operand and result have same element type",
+                 TCresVTEtIsSameAsOpBase<0, 0>>]>,
+    Arguments<(ins StrAttr:$kind, AnyVector:$vector)>,
+    Results<(outs AnyType:$dest)> {
+  let summary = "reduction operation";
+  let description = [{
+    Reduces an 1-D vector "horizontally" into a scalar using the given
+    operation (add/mul/min/max for int/fp and and/or/xor for int only).
+    Note that these operations are restricted to 1-D vectors to remain
+    close to the corresponding LLVM intrinsics:
+
+    http://llvm.org/docs/LangRef.html#experimental-vector-reduction-intrinsics
+
+    Examples:
+    ```
+      %1 = vector.reduction "add", %0 : vector<16xf32> into f32
+
+      %3 = vector.reduction "xor", %2 : vector<4xi32> into i32
+    ```
+  }];
+  let verifier = [{ return ::verify(*this); }];
+  let assemblyFormat = [{
+    $kind `,` $vector attr-dict `:` type($vector) `into` type($dest)
+  }];
+  let extraClassDeclaration = [{
+    VectorType getVectorType() {
+      return vector().getType().cast<VectorType>();
+    }
+  }];
+}
+
+// TODO(ajcbik): quick version with "fused" accumulator; next step
+//               will merge Reduction/ReductionV2 into one with
+//               an optional accumulator instead
+def Vector_ReductionV2Op :
+  Vector_Op<"reductionv2", [NoSideEffect]>,
+    Arguments<(ins StrAttr:$kind, VectorOf<[F32, F64]>:$vector, AnyType:$acc)>,
+    Results<(outs AnyType:$dest)> {
+  let summary = "reduction operation";
+  let description = [{
+     As vector.reduction, but with a fused accumulator (add/mul for fp only).
+  }];
+  let verifier = ?;
+  let assemblyFormat = [{
+    $kind `,` $vector `,` $acc attr-dict `:`
+      type($vector) `,` type($acc) `into` type($dest)
+  }];
+  let extraClassDeclaration = [{
+    VectorType getVectorType() {
+      return vector().getType().cast<VectorType>();
+    }
+  }];
+}
+
 def Vector_BroadcastOp :
   Vector_Op<"broadcast", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
     Arguments<(ins AnyType:$source)>,
     Results<(outs AnyVector:$vector)> {
   let summary = "broadcast operation";
   let description = [{

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

static SmallVector<int64_t, 4> getI64SubArray(ArrayAttr arrayAttr,
   res.reserve(arrayAttr.size() - dropFront - dropBack);
   for (auto it = range.begin() + dropFront, eit = range.end() - dropBack;
        it != eit; ++it)
     res.push_back((*it).getValue().getSExtValue());
   return res;
 }
 
 namespace {
+
 class VectorBroadcastOpConversion : public LLVMOpLowering {
 public:
   explicit VectorBroadcastOpConversion(MLIRContext *context,
                                        LLVMTypeConverter &typeConverter)
       : LLVMOpLowering(vector::BroadcastOp::getOperationName(), context,
                        typeConverter) {}
 
   PatternMatchResult
for (int64_t d = 0; d < dim; ++d) {
       Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter);
       result =
           insertOne(rewriter, lowering, loc, result, expand, llvmType, rank, d);
     }
     return result;
   }
 };
 
+class VectorReductionOpConversion : public LLVMOpLowering {
+public:
+  explicit VectorReductionOpConversion(MLIRContext *context,
+                                       LLVMTypeConverter &typeConverter)
+      : LLVMOpLowering(vector::ReductionOp::getOperationName(), context,
+                       typeConverter) {}
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto reductionOp = cast<vector::ReductionOp>(op);
+    auto kind = reductionOp.kind();
+    Type eltType = reductionOp.dest().getType();
+    Type llvmType = lowering.convertType(eltType);
+    if (eltType.isInteger(32) || eltType.isInteger(64)) {
+      // Integer reductions: add/mul/min/max/and/or/xor.
+      if (kind == "add")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_add>(
+            op, llvmType, operands[0]);
+      else if (kind == "mul")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_mul>(
+            op, llvmType, operands[0]);
+      else if (kind == "min")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smin>(
+            op, llvmType, operands[0]);
+      else if (kind == "max")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_smax>(
+            op, llvmType, operands[0]);
+      else if (kind == "and")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_and>(
+            op, llvmType, operands[0]);
+      else if (kind == "or")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_or>(
+            op, llvmType, operands[0]);
+      else if (kind == "xor")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_xor>(
+            op, llvmType, operands[0]);
+      else
+        return matchFailure();
+      return matchSuccess();
+
+    } else if (eltType.isF32() || eltType.isF64()) {
+      // Floating-point reductions: add/mul/min/max
+      if (kind == "add") {
+        Value zero = rewriter.create<LLVM::ConstantOp>(
+            op->getLoc(), llvmType, rewriter.getZeroAttr(eltType));
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>(
+            op, llvmType, zero, operands[0]);
+      } else if (kind == "mul") {
+        Value one = rewriter.create<LLVM::ConstantOp>(
+            op->getLoc(), llvmType, rewriter.getFloatAttr(eltType, 1.0));
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>(
+            op, llvmType, one, operands[0]);
+      } else if (kind == "min")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>(
+            op, llvmType, operands[0]);
+      else if (kind == "max")
+        rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>(
+            op, llvmType, operands[0]);
+      else
+        return matchFailure();
+      return matchSuccess();
+    }
+    return matchFailure();
+  }
+};
+
+// TODO(ajcbik): merge Reduction and ReductionV2
+class VectorReductionV2OpConversion : public LLVMOpLowering {
+public:
+  explicit VectorReductionV2OpConversion(MLIRContext *context,
+                                         LLVMTypeConverter &typeConverter)
+      : LLVMOpLowering(vector::ReductionV2Op::getOperationName(), context,
+                       typeConverter) {}
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto reductionOp = cast<vector::ReductionV2Op>(op);
+    auto kind = reductionOp.kind();
+    Type eltType = reductionOp.dest().getType();
+    Type llvmType = lowering.convertType(eltType);
+    if (kind == "add") {
+      rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>(
+          op, llvmType, operands[1], operands[0]);
+      return matchSuccess();
+    } else if (kind == "mul") {
+      rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>(
+          op, llvmType, operands[1], operands[0]);
+      return matchSuccess();
+    }
+    return matchFailure();
+  }
+};
+
 class VectorShuffleOpConversion : public LLVMOpLowering {
 public:
   explicit VectorShuffleOpConversion(MLIRContext *context,
                                      LLVMTypeConverter &typeConverter)
       : LLVMOpLowering(vector::ShuffleOp::getOperationName(), context,
                        typeConverter) {}
 
   PatternMatchResult

 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   patterns.insert<VectorFMAOpNDRewritePattern,
                   VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
                   VectorStridedSliceOpConversion>(ctx);
-  patterns.insert<VectorBroadcastOpConversion, VectorShuffleOpConversion,
+  patterns.insert<VectorBroadcastOpConversion, VectorReductionOpConversion,
+                  VectorReductionV2OpConversion, VectorShuffleOpConversion,
                   VectorExtractElementOpConversion, VectorExtractOpConversion,
                   VectorFMAOp1DConversion, VectorInsertElementOpConversion,
                   VectorInsertOpConversion, VectorOuterProductOpConversion,
                   VectorTypeCastOpConversion, VectorPrintOpConversion>(
       ctx, converter);
 }
 
 namespace {
 struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 
 void LowerVectorToLLVMPass::runOnModule() {
-  // Perform progressive lowering of operations on "slices".
-  // Folding and DCE get rid of all non-leaking tuple ops.
+  // Perform progressive lowering of operations on "slices" and
+  // all contraction operations. Also applies folding and DCE.
   {
     OwningRewritePatternList patterns;
     populateVectorSlicesLoweringPatterns(patterns, &getContext());
+    populateVectorContractLoweringPatterns(patterns, &getContext());
     applyPatternsGreedily(getModule(), patterns);
   }
 
   // Convert to the LLVM IR dialect.
   LLVMTypeConverter converter(&getContext());
   OwningRewritePatternList patterns;
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateStdToLLVMConversionPatterns(converter, patterns);

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 }
 
 ArrayAttr vector::getVectorSubscriptAttr(Builder &builder,
                                          ArrayRef<int64_t> values) {
   return builder.getI64ArrayAttr(values);
 }
 
 //===----------------------------------------------------------------------===//
+// ReductionOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verify(ReductionOp op) {
+  // Verify for 1-D vector.
+  int64_t rank = op.getVectorType().getRank();
+  if (rank != 1)
+    return op.emitOpError("unsupported reduction rank: ") << rank;
+
+  // Verify supported reduction kind.
+  auto kind = op.kind();
+  Type eltType = op.dest().getType();
+  if (kind == "add" || kind == "mul" || kind == "min" || kind == "max") {
+    if (eltType.isF32() || eltType.isF64() || eltType.isInteger(32) ||
+        eltType.isInteger(64))
+      return success();
+    return op.emitOpError("unsupported reduction type");
+  }
+  if (kind == "and" || kind == "or" || kind == "xor") {
+    if (eltType.isInteger(32) || eltType.isInteger(64))
+      return success();
+    return op.emitOpError("unsupported reduction type");
+  }
+  return op.emitOpError("unknown reduction kind: ") << kind;
+}
+
+//===----------------------------------------------------------------------===//
 // ContractionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ContractionOp::build(Builder *builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
                                   ArrayRef<StringRef> iteratorTypes) {
   result.addOperands({lhs, rhs, acc});
void ContractionOp::getIterationBounds(
   SmallVector<int64_t, 2> iterationShape;
   for (auto it : llvm::enumerate(iterator_types())) {
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), getContext());
     auto iteratorTypeName = it.value().cast<StringAttr>().getValue();
     if (iteratorTypeName == getReductionIteratorTypeName()) {
       // Get reduction dim size from lhs shape (same size in rhsShape).
       int64_t lhsDimIndex = getResultIndex(indexingMaps[0], targetExpr);
+      if (lhsDimIndex < 0) {
+        continue;
+      }
       assert(lhsDimIndex >= 0);
       iterationBounds.push_back(lhsShape[lhsDimIndex]);
       continue;
     }
     // Get parallel dimension size from result shape.
     int64_t resDimIndex = getResultIndex(indexingMaps[2], targetExpr);
+    if (resDimIndex < 0) {
+      continue;
+    }
     assert(resDimIndex >= 0);
     assert(resVectorType != nullptr);
     iterationBounds.push_back(resVectorType.getShape()[resDimIndex]);
   }
 }
 
 void ContractionOp::getIterationIndexMap(
     std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap) {

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

for (int64_t i = 0; i < tupleSize; ++i) {
           loc, tupleGet, result, elementOffsets, strides);
     }
 
     rewriter.replaceOp(op, result);
     return matchSuccess();
   }
 };
 
+/// Progressive lowering of ConstractionOp.
+class ContractionOpLowering : public OpRewritePattern<vector::ContractionOp> {
+public:
+  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(vector::ContractionOp op,
+                                     PatternRewriter &rewriter) const override {
+    // TODO(ajcbik): implement masks
+    if (llvm::size(op.masks()) != 0)
+      return matchFailure();
+
+    auto loc = op.getLoc();
+    VectorType lhsType = op.getLhsType();
+    VectorType rhsType = op.getRhsType();
+    Type resType = op.getResultType();
+
+    // Find first batch dimension in lhs/rhs, and lower when found.
+    std::vector<std::pair<int64_t, int64_t>> batchDimMap = op.getBatchDimMap();
+    if (!batchDimMap.empty()) {
+      // TODO(ajcbik): implement batch
+      return matchFailure();
+    }
+
+    // Collect contracting dimensions.
+    std::vector<std::pair<int64_t, int64_t>> contractingDimMap =
+        op.getContractingDimMap();
+    DenseSet<int64_t> lhsContractingDimSet;
+    DenseSet<int64_t> rhsContractingDimSet;
+    for (auto &dimPair : contractingDimMap) {
+      lhsContractingDimSet.insert(dimPair.first);
+      rhsContractingDimSet.insert(dimPair.second);
+    }
+
+    // Find free dimension in lhs/rhs, and lower first when found.
+    for (int64_t i = 0, e = lhsType.getRank(); i < e; ++i) {
+      if (lhsContractingDimSet.count(i) == 0) {
+        // TODO(ajcbik): implement free
+        return matchFailure();
+      }
+    }
+    for (int64_t i = 0, e = rhsType.getRank(); i < e; ++i) {
+      if (rhsContractingDimSet.count(i) == 0) {
+        // TODO(ajcbik): implement free
+        return matchFailure();
+      }
+    }
+
+    // Only contraction dimensions remain.
+    if (!resType.isa<VectorType>() && lhsType.getRank() == 1 &&
+        rhsType.getRank() == 1) {
+      // Handle reduction into scalar.
+      Value zero = rewriter.create<ConstantOp>(loc, resType,
+                                               rewriter.getZeroAttr(resType));
+      Value splat = rewriter.create<SplatOp>(loc, lhsType, zero);
+      Value fma =
+          rewriter.create<vector::FMAOp>(loc, op.lhs(), op.rhs(), splat);
+      StringAttr kind = rewriter.getStringAttr("add");
+      rewriter.replaceOpWithNewOp<vector::ReductionV2Op>(op, resType, kind, fma,
+                                                         op.acc());
+      return matchSuccess();
+    }
+    // TODO(ajcbik): implement more contraction
+    return matchFailure();
+  }
+};
+
 } // namespace
 
 // TODO(andydavis) Add pattern to rewrite ExtractSlices(ConstantMaskOp).
 // TODO(andydavis) Add this as DRR pattern.
 void mlir::vector::populateVectorToVectorTransformationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<SplitTransferReadOp, SplitTransferWriteOp, TupleGetFolderOp>(
       context);
 }
 
 void mlir::vector::populateVectorSlicesLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<ExtractSlicesOpLowering, InsertSlicesOpLowering>(context);
 }
+
+void mlir::vector::populateVectorContractLoweringPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<ContractionOpLowering>(context);
+}

mlir/test/Dialect/VectorOps/vector-contract-transforms.mlir

+// RUN: mlir-opt %s -test-vector-contraction-conversion | FileCheck %s
+
+#dotp_accesses = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> ()>
+]
+#dotp_trait = {
+  indexing_maps = #dotp_accesses,
+  iterator_types = ["reduction"]
+}
+
+// CHECK-LABEL: func @extract_contract1
+// CHECK-SAME: %[[A:.*0]]: vector<4xf32>,
+// CHECK-SAME: %[[B:.*1]]: vector<4xf32>,
+// CHECK-SAME: %[[C:.*2]]: f32
+// CHECK:      %[[Z:.*]] = constant dense<0.000000e+00>
+// CHECK:      %[[F:.*]] = vector.fma %[[A]], %[[B]], %[[Z]] : vector<4xf32>
+// CHECK:      %[[R:.*]] = vector.reductionv2 "add", %[[F]], %[[C]] : vector<4xf32>, f32 into f32
+// CHECK:      return %[[R]]
+
+func @extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: f32) -> f32 {
+  %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2 : vector<4xf32>, vector<4xf32> into f32
+  return %0 : f32
+}

mlir/test/lib/Transforms/TestVectorTransforms.cpp

struct TestVectorSlicesConversion
     : public FunctionPass<TestVectorSlicesConversion> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     populateVectorSlicesLoweringPatterns(patterns, &getContext());
     applyPatternsGreedily(getFunction(), patterns);
   }
 };
 
+struct TestVectorContractionConversion
+    : public FunctionPass<TestVectorContractionConversion> {
+  void runOnFunction() override {
+    OwningRewritePatternList patterns;
+    populateVectorContractLoweringPatterns(patterns, &getContext());
+    applyPatternsGreedily(getFunction(), patterns);
+  }
+};
+
 } // end anonymous namespace
 
-static PassRegistration<TestVectorToVectorConversion>
-    pass("test-vector-to-vector-conversion",
-         "Test conversion patterns between ops in the vector dialect");
+static PassRegistration<TestVectorToVectorConversion> vector_to_vector_pass(
+    "test-vector-to-vector-conversion",
 
 static PassRegistration<TestVectorSlicesConversion> slices_pass(
     "test-vector-slices-conversion",
     "Test conversion patterns that lower slices ops in the vector dialect");
+
+static PassRegistration<TestVectorContractionConversion>
+    contraction_pass("test-vector-contraction-conversion",
+                     "Test conversion patterns that lower contraction ops in "
+                     "the vector dialect");