[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 ↗

(On Diff #244320)

trivial braces

373 ↗

(On Diff #244320)

trivial braces

mlir/test/lib/Transforms/TestVectorTransforms.cpp

55–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

def Vector_ReductionOp :
   }];
   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

if (eltType.isInteger(32) || eltType.isInteger(64)) {
       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

 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   patterns.insert<VectorFMAOpNDRewritePattern,
                   VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
                   VectorStridedSliceOpConversion>(ctx);
   patterns.insert<VectorBroadcastOpConversion, VectorReductionOpConversion,
-                  VectorShuffleOpConversion, VectorExtractElementOpConversion,
-                  VectorExtractOpConversion, VectorFMAOp1DConversion,
-                  VectorInsertElementOpConversion, VectorInsertOpConversion,
-                  VectorOuterProductOpConversion, VectorTypeCastOpConversion,
-                  VectorPrintOpConversion>(ctx, converter);
+                  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/VectorTransforms.cpp

for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
     if (lastPos.hasValue() && currPos != lastPos.getValue() + 1)
       return false;
     lastPos = currPos;
   }
   return true;
 }
 
 namespace {
+
 // Splits vector TransferReadOp into smaller TransferReadOps based on slicing
 // scheme of its unique ExtractSlicesOp user.
 struct SplitTransferReadOp : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(vector::TransferReadOp xferReadOp,
                                      PatternRewriter &rewriter) const override {
     // TODO(andydavis, ntv) Support splitting TransferReadOp with non-identity
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<ShapeCastOpDecomposer, ShapeCastOpFolder, 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]]
+// CHECK:      return %[[R]] : f32
+
+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
 
 namespace mlir {
 void registerTestVectorConversions() {
-  PassRegistration<TestVectorToVectorConversion> pass(
+  PassRegistration<TestVectorToVectorConversion> vectorToVectorPass(
       "test-vector-to-vector-conversion",
       "Test conversion patterns between ops in the vector dialect");
 
-  PassRegistration<TestVectorSlicesConversion> slices_pass(
+  PassRegistration<TestVectorSlicesConversion> slicesPass(
       "test-vector-slices-conversion",
       "Test conversion patterns that lower slices ops in the vector dialect");
+
+  PassRegistration<TestVectorContractionConversion> contractionPass(
+      "test-vector-contraction-conversion",
+      "Test conversion patterns that lower contract ops in the vector dialect");
 }
 } // namespace mlir