[mlir][Vector] Mostly-NFC - Restructure options for lowering to LLVM Matrix Intrinsics

[nicolasvasilache]

Summary

This revision restructures the calling of vector transforms to make it more flexible to ask for lowering through LLVM matrix intrinsics.
This also makes sure we bail out in degenerate cases (i.e. 1) in which LLVM complains about not being able to scalarize.

[mehdi_amini]

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

955 ↗

(On Diff #250687)

Can you expand the comment to describe the problem with LLVM?
Also is it a fundamental problem is it just a temporary bug/limitation in LLVM that should be lifted there? (what is the tracking bug? Is there a test-case?)

mlir/test/lib/Transforms/TestVectorTransforms.cpp

27

Can you use pass options instead?

[aartbik]

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

27 ↗

(On Diff #250687)

add a comment here to document the struct

(options for the vector transforms, or at least the contract populate method for now)

28 ↗

(On Diff #250687)

provide convenience constructors for this

(one bool is simple but if we add more that may come in handy?)

[nicolasvasilache]

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

28 ↗

(On Diff #250687)

What benefit do constructors bring in this instance?
A simple struct with public members is the simplest thing in my book.

Once multiple options are available and they can be logically grouped and given a meaningful name we can provide static constructors that configure a set of options.

We can also provide chaining semantics to have something like: VectorTransformsOptions().setX(true).setY(3)

All these will make sense in the future but are premature until we have at least 3-5 flags + behaviors that must compound.

mlir/test/lib/Transforms/TestVectorTransforms.cpp

27

very cool! Done.

[nicolasvasilache]

Address review comments.

[aartbik]

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

28 ↗

(On Diff #250687)

Well, this was just a vague forward looking idea; with proper enums a constructor could read something like this

VectorTransformOpions(ENABLE_LOWER_TO_LLVM_MATRIX, USE_THIS_AND_THAT)

when passing it into a method without the need for a chain of setters. But given that we have one bool, the point is moot :-)

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

956 ↗

(On Diff #250849)

if this is really a for now, consider adding a TODO in addition to the bug links

[nicolasvasilache]

Add TODO.

[nicolasvasilache]

Rebase

mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h

 void populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
 
 /// Collect a set of patterns to convert from the Vector dialect to LLVM.
 void populateVectorToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                             OwningRewritePatternList &patterns);
 
 /// Create a pass to convert vector operations to the LLVMIR dialect.
-OpPassBase<ModuleOp> *createLowerVectorToLLVMPass();
+std::unique_ptr<OpPassBase<ModuleOp>> createConvertVectorToLLVMPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLLVM_H_

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

 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Interfaces/SideEffects.h"
 
 namespace mlir {
 class MLIRContext;
 class OwningRewritePatternList;
 namespace vector {
 
+/// Structure to control the behavior of vector transform patterns.
+struct VectorTransformsOptions {
+  /// Let vector.contract lower to vector.matrix_multiply and LLVM matrix
+  /// intrinsics.
+  bool lowerToLLVMMatrixIntrinsics = false;
+};
+
 /// Collect a set of vector-to-vector canonicalization patterns.
 void populateVectorToVectorCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context);
 
 /// Collect a set of vector-to-vector transformation patterns.
 void populateVectorToVectorTransformationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context);
 

 
 /// Collect a set of transformation patterns that are related to contracting
 /// or expanding vector operations:
 ///   ContractionOpLowering,
 ///   ShapeCastOp2DDownCastRewritePattern, ShapeCastOp2DUpCastRewritePattern
 ///   OuterproductOpLowering
 /// These transformation express higher level vector ops in terms of more
 /// elementary extraction, insertion, reduction, product, and broadcast ops.
-void populateVectorContractLoweringPatterns(OwningRewritePatternList &patterns,
-                                            MLIRContext *context);
+void populateVectorContractLoweringPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context,
+    VectorTransformsOptions vectorTransformOptions = VectorTransformsOptions());
 
 /// 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);
 

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

void ConvertLinalgToLLVMPass::runOnModule() {
 
   // Convert to the LLVM IR dialect using the converter defined above.
   OwningRewritePatternList patterns;
   LLVMTypeConverter converter(&getContext());
   populateAffineToStdConversionPatterns(patterns, &getContext());
   populateLoopToStdConversionPatterns(patterns, &getContext());
   populateStdToLLVMConversionPatterns(converter, patterns, /*useAlloca=*/false,
                                       /*emitCWrappers=*/true);
+  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateLinalgToStandardConversionPatterns(patterns, &getContext());
   populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext());
 
   LLVMConversionTarget target(getContext());
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   target.addLegalOp<ModuleOp, ModuleTerminatorOp>();

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

void LowerVectorToLLVMPass::runOnModule() {
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   if (failed(
           applyPartialConversion(getModule(), target, patterns, &converter))) {
     signalPassFailure();
   }
 }
 
-OpPassBase<ModuleOp> *mlir::createLowerVectorToLLVMPass() {
-  return new LowerVectorToLLVMPass();
+std::unique_ptr<OpPassBase<ModuleOp>> mlir::createConvertVectorToLLVMPass() {
+  return std::make_unique<LowerVectorToLLVMPass>();
 }
 
 static PassRegistration<LowerVectorToLLVMPass>
     pass("convert-vector-to-llvm",
          "Lower the operations from the vector dialect into the LLVM dialect");

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 
 using namespace mlir;
 using llvm::dbgs;
 using mlir::functional::zipMap;
 
-static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
-
-static llvm::cl::opt<bool> lowerToLLVMMatrixIntrinsics(
-    "vector-lower-matrix-intrinsics",
-    llvm::cl::desc("Lower vector.contract to llvm.intr.matrix.multiply"),
-    llvm::cl::init(false), llvm::cl::cat(clOptionsCategory));
-
 /// Given a shape with sizes greater than 0 along all dimensions,
 /// returns the distance, in number of elements, between a slice in a dimension
 /// and the next slice in the same dimension.
 ///   e.g. shape[3, 4, 5] -> linearization_basis[20, 5, 1]
 static SmallVector<int64_t, 8> computeStrides(ArrayRef<int64_t> shape) {
   if (shape.empty())
     return {};
   SmallVector<int64_t, 8> tmp;

 ///
 /// TODO(ajcbik): break down into transpose/reshape/cast ops
 ///               when they become available to avoid code dup
 /// TODO(ajcbik): investigate lowering order impact on performance
 class ContractionOpLowering : public OpRewritePattern<vector::ContractionOp> {
 public:
   using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
 
+  ContractionOpLowering(vector::VectorTransformsOptions vectorTransformsOptions,
+                        MLIRContext *context)
+      : OpRewritePattern<vector::ContractionOp>(context),
+        vectorTransformsOptions(vectorTransformsOptions) {}
+
   PatternMatchResult matchAndRewrite(vector::ContractionOp op,
                                      PatternRewriter &rewriter) const override {
     // TODO(ajcbik): implement masks
     if (llvm::size(op.masks()) != 0)
       return matchFailure();
 
     // TODO(ntv, ajcbik): implement benefits, cost models, separate this out in
     // a new pattern.
     // TODO(ntv, fhahn): once row-major mode is available in LLVM's matrix
     // intrinsics, use that.
-    if (lowerToLLVMMatrixIntrinsics &&
+    if (vectorTransformsOptions.lowerToLLVMMatrixIntrinsics &&
         isColumnMajorMatmul(op.indexing_maps())) {
       VectorType lhsType = op.getLhsType();
       VectorType rhsType = op.getRhsType();
+      unsigned lhsRows = op.getLhsType().getShape()[0];
+      unsigned lhsColumns = op.getLhsType().getShape()[1];
+      unsigned rhsColumns = op.getRhsType().getShape()[1];
+
+      // In cases where matrices are degenerate, scalarization issues occur in
+      // the backend. Avoid all LLVM scalarization issues for now.
+      // For more details, see: https://bugs.llvm.org/show_bug.cgi?id=45227 and
+      // https://bugs.llvm.org/show_bug.cgi?id=45229
+      // TODO(ntv, fhahn): Relax once above bugs are fixed.
+      if (lhsRows != 1 && lhsColumns != 1 && rhsColumns != 1) {
         Type flattenedLHSType =
             VectorType::get(lhsType.getNumElements(), lhsType.getElementType());
         Type flattenedRHSType =
             VectorType::get(rhsType.getNumElements(), rhsType.getElementType());
         auto lhs = rewriter.create<vector::ShapeCastOp>(
             op.getLoc(), flattenedLHSType, op.lhs());
         auto rhs = rewriter.create<vector::ShapeCastOp>(
             op.getLoc(), flattenedRHSType, op.rhs());
 
-      unsigned lhsRows = op.getLhsType().getShape()[0];
-      unsigned lhsColumns = op.getLhsType().getShape()[1];
-      unsigned rhsColumns = op.getRhsType().getShape()[1];
         Value mul = rewriter.create<vector::MatmulOp>(
             op.getLoc(), lhs, rhs, lhsRows, lhsColumns, rhsColumns);
         mul = rewriter.create<vector::ShapeCastOp>(op.getLoc(),
                                                    op.acc().getType(), mul);
         Type elementType = op.getLhsType().getElementType();
         assert(elementType.isIntOrFloat());
         if (elementType.isa<IntegerType>())
           rewriter.replaceOpWithNewOp<AddIOp>(op, op.acc(), mul);
         else
           rewriter.replaceOpWithNewOp<AddFOp>(op, op.acc(), mul);
         return matchSuccess();
       }
+    }
 
     // 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()) {
       int64_t lhsIndex = batchDimMap[0].first;
       int64_t rhsIndex = batchDimMap[0].second;
       rewriter.replaceOp(op, lowerParallel(op, lhsIndex, rhsIndex, rewriter));
       return matchSuccess();
for (int64_t d = 0, e = type.getDimSize(0); d < e; d++) {
       Value ins =
           rewriter.create<vector::ExtractOp>(loc, insType, val, posAttr);
       Value sto = reshapeStore(loc, ins, ext, vType, index - 1, pos, rewriter);
       result =
           rewriter.create<vector::InsertOp>(loc, type, sto, result, posAttr);
     }
     return result;
   }
+
+  vector::VectorTransformsOptions vectorTransformsOptions;
 };
 
 /// ShapeOp 2D -> 1D downcast serves the purpose of flattening 2-D to 1-D
 /// vectors progressively on the way to target llvm.matrix intrinsics.
 /// This iterates over the most major dimension of the 2-D vector and performs
 /// rewrites into:
 ///   vector.extract from 2-D + vector.insert_strided_slice offset into 1-D
 class ShapeCastOp2DDownCastRewritePattern

 }
 
 void mlir::vector::populateVectorSlicesLoweringPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<ExtractSlicesOpLowering, InsertSlicesOpLowering>(context);
 }
 
 void mlir::vector::populateVectorContractLoweringPatterns(
-    OwningRewritePatternList &patterns, MLIRContext *context) {
-  patterns.insert<ContractionOpLowering, ShapeCastOp2DDownCastRewritePattern,
+    OwningRewritePatternList &patterns, MLIRContext *context,
+    VectorTransformsOptions parameters) {
+  patterns.insert<ShapeCastOp2DDownCastRewritePattern,
                   ShapeCastOp2DUpCastRewritePattern, OuterProductOpLowering>(
       context);
+  patterns.insert<ContractionOpLowering>(parameters, context);
 }

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

 // RUN: mlir-opt %s -test-vector-contraction-conversion | FileCheck %s
-// RUN: mlir-opt %s -test-vector-contraction-conversion -vector-lower-matrix-intrinsics | FileCheck %s --check-prefix=MATRIX
+// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-lower-matrix-intrinsics=1 | FileCheck %s --check-prefix=MATRIX
 
 #dotp_accesses = [
   affine_map<(i) -> (i)>,
   affine_map<(i) -> (i)>,
   affine_map<(i) -> ()>
 ]
 #dotp_trait = {
   indexing_maps = #dotp_accesses,

mlir/test/lib/Transforms/TestVectorTransforms.cpp

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorTransforms.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 using namespace mlir::vector;
-
 namespace {
 
 #include "TestVectorTransformPatterns.h.inc"
 
 struct TestVectorToVectorConversion
     : public FunctionPass<TestVectorToVectorConversion> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     auto *context = &getContext();
     populateWithGenerated(context, &patterns);
     populateVectorToVectorCanonicalizationPatterns(patterns, context);
     populateVectorToVectorTransformationPatterns(patterns, context);
     applyPatternsGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorSlicesConversion
     : public FunctionPass<TestVectorSlicesConversion> {
   void runOnFunction() override {
     OwningRewritePatternList patterns;
     populateVectorSlicesLoweringPatterns(patterns, &getContext());
     applyPatternsGreedily(getFunction(), patterns);
   }
 };
 
 struct TestVectorContractionConversion
     : public FunctionPass<TestVectorContractionConversion> {
+  TestVectorContractionConversion() = default;
+  TestVectorContractionConversion(const TestVectorContractionConversion &pass) {
+  }
+
+  Option<bool> lowerToLLVMMatrixIntrinsics{
+      *this, "vector-lower-matrix-intrinsics",
+      llvm::cl::desc("Lower vector.contract to llvm.intr.matrix.multiply"),
+      llvm::cl::init(false)};
+
   void runOnFunction() override {
     OwningRewritePatternList patterns;
-    populateVectorContractLoweringPatterns(patterns, &getContext());
+    VectorTransformsOptions options{
+        /*lowerToLLVMMatrixIntrinsics=*/lowerToLLVMMatrixIntrinsics};
+    populateVectorContractLoweringPatterns(patterns, &getContext(), options);
     applyPatternsGreedily(getFunction(), patterns);
   }
 };
 
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestVectorConversions() {