[mlir][Vector] Add a vector.matrix_multiply op on 1-D vectors

[nicolasvasilache]

Summary

This op mirrorrs the llvm.intr counterpart and allows lowering + type conversions in a progressive fashion.

[aartbik]

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

20

why is this in its own population (instead of adding it to the general VectorToLLVM)?
I don't object, but just curious. I can see isolated testing, but you don't seem to exploit that

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

1343

Some parts of this comments could be moved into description (e.g. for some others, we also say why we picked it, like being close to the llvm implementation)

1362

can we get a summary and description please? we need to keep our docs a bit up to date ;-)

[nicolasvasilache]

Address review comments, fix incorrect rhs_columns attribute name.

[nicolasvasilache]

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

20

Since lowering matrix intrinsics in LLVM requires a special flag, it makes sense to group those patterns together.
Updated the doc to reflect this.

[nicolasvasilache]

Update comments.

[mehdi_amini]

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

20

(Typo: missing space in wheninvoking)

Why do we need a special flag for matrix lowering?

[nicolasvasilache]

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

20

Not sure, @fhahn ?

[nicolasvasilache]

Typo.

[aartbik]

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

20

typo: wheninvoking -> when invoking

mlir/include/mlir/IR/OpBase.td

490 ↗

(On Diff #248934)

rebase this

[nicolasvasilache]

Rebase.

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

 
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 class LLVMTypeConverter;
 class ModuleOp;
 template <typename T> class OpPassBase;
 
+/// Collect a set of patterns to convert from Vector contractions to LLVM Matrix
+/// Intrinsics. To lower to assembly, the LLVM flag -lower-matrix-intrinsics
+/// will be needed when invoking LLVM.
+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();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLLVM_H_

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 }
 
 /// Create a llvm.matrix.multiply call, multiplying 2-D matrices LHS and RHS, as
 /// specified in the LLVM MatrixBuilder.
 def LLVM_MatrixMultiplyOp
     : LLVM_OneResultOp<"intr.matrix.multiply">,
       Arguments<(
         ins LLVM_Type:$lhs, LLVM_Type:$rhs,
-            I32Attr:$lhs_rows, I32Attr:$lhs_columns, I32Attr:$rhs_rows)> {
+            I32Attr:$lhs_rows, I32Attr:$lhs_columns, I32Attr:$rhs_columns)> {
   string llvmBuilder = [{
     llvm::MatrixBuilder<decltype(builder)> mb(builder);
     $res = mb.CreateMatrixMultiply(
       $lhs, $rhs, $lhs_rows.getZExtValue(), $lhs_columns.getZExtValue(),
-      $rhs_rows.getZExtValue());
+      $rhs_columns.getZExtValue());
   }];
   let assemblyFormat = "$lhs `,` $rhs attr-dict "
     "`:` `(` type($lhs) `,` type($rhs) `)` `->` type($res)";
 }
 
 /// Create a llvm.matrix.transpose call, transposing a `rows` x `columns` 2-D
 /// `matrix`, as specified in the LLVM MatrixBuilder.
 def LLVM_MatrixTranposeOp

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

def Vector_PrintOp :
   let extraClassDeclaration = [{
     Type getPrintType() {
       return source().getType();
     }
   }];
   let assemblyFormat = "$source attr-dict `:` type($source)";
 }
 
+//===----------------------------------------------------------------------===//
+// Ops used for supporting progressive lowering and conversion type changes.
+//===----------------------------------------------------------------------===//
+
+/// Vector dialect matrix multiplication op that operates on flattened 1-D
+/// MLIR vectors. This is the counterpart of llvm.matrix.multiply in MLIR.
+/// This may seem redundant with vector.contract but it serves the purposes of
+/// more progressive lowering and localized type conversion on the path:
+///   `vector<...x...xf32> -> vector<...xf32> -> !llvm<... x float>`.
+def Vector_MatmulOp : Vector_Op<"matrix_multiply", [NoSideEffect,
+        PredOpTrait<"lhs operand and result have same element type",
+                    TCresVTEtIsSameAsOpBase<0, 0>>,
+        PredOpTrait<"rhs operand and result have same element type",
+                    TCresVTEtIsSameAsOpBase<0, 1>>]>,
+      Arguments<(
+        // TODO(ntv, fhahn): tighten vector element types that make sense.
+        ins VectorOfRankAndType<[1],
+              [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$lhs,
+            VectorOfRankAndType<[1],
+              [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$rhs,
+            I32Attr:$lhs_rows, I32Attr:$lhs_columns, I32Attr:$rhs_columns)>,
+      Results<(
+        outs VectorOfRankAndType<[1],
+               [AnySignlessInteger, AnySignedInteger, AnyFloat]>:$res)>
+{
+  let summary = "Vector matrix multiplication op that operates on flattened 1-D"
+    " MLIR vectors";
+  let description = [{
+    This is the counterpart of llvm.matrix.multiply in MLIR. It serves the
+    purposes of more progressive lowering and localized type conversion.
+
+    The ‘vector.matrix_multiply’ op treats `lhs` as matrix with <lhs_rows> rows
+    and <lhs_columns> columns, `rhs` as matrix with <lhs_columns> rows and
+    <rhs_columns> and multiplies them. The result matrix is returned embedded in
+    the result vector.
+
+    Example:
+
+    ```
+      %C = vector.matrix_multiply %A, %B
+        { lhs_rows = 4: i32, lhs_columns = 16: i32 , rhs_columns = 3: i32 } :
+        (vector<64xf64>, vector<48xf64>) -> vector<12xf64>
+    ```
+  }];
+  let builders = [
+   OpBuilder<"Builder *builder, OperationState &result, Value lhs, Value rhs, "
+             "unsigned lhsRows, unsigned lhsColumns, unsigned rhsColumns",
+   [{
+     result.addOperands({lhs, rhs});
+     result.addAttribute("lhs_rows", builder->getI32IntegerAttr(lhsRows));
+     result.addAttribute("lhs_columns", builder->getI32IntegerAttr(lhsColumns));
+     result.addAttribute("rhs_columns", builder->getI32IntegerAttr(rhsColumns));
+     result.addTypes(VectorType::get(lhsRows * lhsColumns,
+       lhs.getType().cast<VectorType>().getElementType()));
+   }]>,
+  ];
+  let verifier = ?;
+  let assemblyFormat = "$lhs `,` $rhs attr-dict "
+    "`:` `(` type($lhs) `,` type($rhs) `)` `->` type($res)";
+}
+
 #endif // VECTOR_OPS

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

for (int64_t d = 0; d < dim; ++d) {
       Value expand = expandRanks(one, loc, redSrcType, redDstType, rewriter);
       result = insertOne(rewriter, typeConverter, loc, result, expand, llvmType,
                          rank, d);
     }
     return result;
   }
 };
 
+/// Conversion pattern for a vector.matrix_multiply.
+/// This is lowered directly to the proper llvm.intr.matrix.multiply.
+class VectorMatmulOpConversion : public ConvertToLLVMPattern {
+public:
+  explicit VectorMatmulOpConversion(MLIRContext *context,
+                                    LLVMTypeConverter &typeConverter)
+      : ConvertToLLVMPattern(vector::MatmulOp::getOperationName(), context,
+                             typeConverter) {}
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto matmulOp = cast<vector::MatmulOp>(op);
+    auto adaptor = vector::MatmulOpOperandAdaptor(operands);
+    rewriter.replaceOpWithNewOp<LLVM::MatrixMultiplyOp>(
+        op, typeConverter.convertType(matmulOp.res().getType()), adaptor.lhs(),
+        adaptor.rhs(), matmulOp.lhs_rows(), matmulOp.lhs_columns(),
+        matmulOp.rhs_columns());
+    return matchSuccess();
+  }
+};
+
 class VectorReductionOpConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorReductionOpConversion(MLIRContext *context,
                                        LLVMTypeConverter &typeConverter)
       : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context,
                              typeConverter) {}
 
   PatternMatchResult
void mlir::populateVectorToLLVMConversionPatterns(
   patterns.insert<VectorBroadcastOpConversion, VectorReductionOpConversion,
                   VectorShuffleOpConversion, VectorExtractElementOpConversion,
                   VectorExtractOpConversion, VectorFMAOp1DConversion,
                   VectorInsertElementOpConversion, VectorInsertOpConversion,
                   VectorOuterProductOpConversion, VectorTypeCastOpConversion,
                   VectorPrintOpConversion>(ctx, converter);
 }
 
+void mlir::populateVectorToLLVMMatrixConversionPatterns(
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
+  MLIRContext *ctx = converter.getDialect()->getContext();
+  patterns.insert<VectorMatmulOpConversion>(ctx, converter);
+}
+
 namespace {
 struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 
 void LowerVectorToLLVMPass::runOnModule() {
   // 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;
+  populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateVectorToLLVMConversionPatterns(converter, patterns);
   populateStdToLLVMConversionPatterns(converter, patterns);
 
   LLVMConversionTarget target(getContext());
   target.addDynamicallyLegalOp<FuncOp>(
       [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
   if (failed(
           applyPartialConversion(getModule(), target, patterns, &converter))) {

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

func @reduce_i64(%arg0: vector<16xi64>) -> i64 {
   %0 = vector.reduction "add", %arg0 : vector<16xi64> into i64
   return %0 : i64
 }
 // CHECK-LABEL: llvm.func @reduce_i64
 // CHECK-SAME: %[[A:.*]]: !llvm<"<16 x i64>">
 //      CHECK: %[[V:.*]] = "llvm.intr.experimental.vector.reduce.add"(%[[A]])
 //      CHECK: llvm.return %[[V]] : !llvm.i64
 
+
+//                          4x16                16x3               4x3
+func @matrix_ops(%A: vector<64xf64>, %B: vector<48xf64>) -> vector<12xf64> {
+  %C = vector.matrix_multiply %A, %B
+    { lhs_rows = 4: i32, lhs_columns = 16: i32 , rhs_columns = 3: i32 } :
+    (vector<64xf64>, vector<48xf64>) -> vector<12xf64>
+  return %C: vector<12xf64>
+}
+// CHECK-LABEL: llvm.func @matrix_ops
+//       CHECK:   llvm.intr.matrix.multiply %{{.*}}, %{{.*}} {
+//  CHECK-SAME: lhs_columns = 16 : i32, lhs_rows = 4 : i32, rhs_columns = 3 : i32
+//  CHECK-SAME: } : (!llvm<"<64 x double>">, !llvm<"<48 x double>">) -> !llvm<"<12 x double>">

mlir/test/Target/llvmir-intrinsics.mlir

 }
 
 // CHECK-LABEL: @matrix_intrinsics
 //                                       4x16                       16x3
 llvm.func @matrix_intrinsics(%A: !llvm<"<64 x float>">, %B: !llvm<"<48 x float>">,
                              %ptr: !llvm<"float*">, %stride: !llvm.i32) {
   // CHECK: call <12 x float> @llvm.matrix.multiply.v12f32.v64f32.v48f32(<64 x float> %0, <48 x float> %1, i32 4, i32 16, i32 3)
   %C = llvm.intr.matrix.multiply %A, %B
-    { lhs_rows = 4: i32, lhs_columns = 16: i32 , rhs_rows = 3: i32} :
+    { lhs_rows = 4: i32, lhs_columns = 16: i32 , rhs_columns = 3: i32} :
     (!llvm<"<64 x float>">, !llvm<"<48 x float>">) -> !llvm<"<12 x float>">
   // CHECK: call <48 x float> @llvm.matrix.transpose.v48f32(<48 x float> %1, i32 3, i32 16)
   %D = llvm.intr.matrix.transpose %B { rows = 3: i32, columns = 16: i32} :
     !llvm<"<48 x float>"> into !llvm<"<48 x float>">
   // CHECK: call <48 x float> @llvm.matrix.columnwise.load.v48f32.p0f32(float* %2, i32 %3, i32 3, i32 16)
   %E = llvm.intr.matrix.columnwise.load %ptr, <stride=%stride>
     { rows = 3: i32, columns = 16: i32} :
     !llvm<"<48 x float>"> from !llvm<"float*"> stride !llvm.i32