[MLIR] Add Linalg support for integer (generalized) matmuls

[GMNGeoffrey]

Summary

This patch adds Linalg named ops for various types of integer matmuls.
Due to limitations in the tc spec/linalg-ods-gen ops cannot be type
polymorphic, so this instead creates new ops (improvements to the
methods for defining Linalg named ops are underway with a prototype at
https://github.com/stellaraccident/mlir-linalgpy).

To avoid the necessity of directly referencing these many new ops, this
adds additional methods to ContractionOpInterface to allow classifying
types of operations based on their indexing maps.

[mravishankar]

Looks fine to me, but would suggest waiting for @nicolasvasilache

[bondhugula]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

40

Typo: the -> that

51

Likewise and everywhere below.

[GMNGeoffrey]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

40

No "the" is correct here: "the given op"

[bondhugula]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

40

There are two "the"s in the sentence - I was referring to the second one.

the correspond -> that correspond

[GMNGeoffrey]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

40

Ah, quite right, thank you :-) https://reviews.llvm.org/D97220 is the fix

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

def LinalgContractionOpInterface : OpInterface<"ContractionOpInterface"> {
    As a consequence, when vectorization of such an op occurs, the only special
    behavior is that the (unique) MulOpType is vectorized into a
    `vector.contract`. All other ops are handled in a generic fashion.
    In the future, we may wish to allow more input arguments and elementwise and
    constant operations that do not involve the reduction dimension(s).
   }];
   let cppNamespace = "::mlir::linalg";
   let verify = [{ return detail::verifyContractionInterface($_op); }];
+  let methods = [
+    InterfaceMethod<
+    /*desc=*/[{
+      Returns whether the given op has indexing maps the correspond to a
+      row-major matmul operation.
+    }],
+    /*retTy=*/"bool",
+    /*methodName=*/"isRowMajorMatmul",
+    /*args=*/(ins),
+    /*methodBody=*/[{
+        return mlir::isRowMajorMatmul($_op.indexing_maps());
+    }]>,
+    InterfaceMethod<
+    /*desc=*/[{
+      Returns whether the given op has indexing maps the correspond to a
+      column-major matmul operation.
+    }],
+    /*retTy=*/"bool",
+    /*methodName=*/"isColumnMajorMatmul",
+    /*args=*/(ins),
+    /*methodBody=*/[{
+        return mlir::isColumnMajorMatmul($_op.indexing_maps());
+    }]>,
+    InterfaceMethod<
+    /*desc=*/[{
+      Returns whether the given op has indexing maps the correspond to a
+      row-major batch matmul operation.
+    }],
+    /*retTy=*/"bool",
+    /*methodName=*/"isRowMajorBatchMatmul",
+    /*args=*/(ins),
+    /*methodBody=*/[{
+        return mlir::isRowMajorBatchMatmul($_op.indexing_maps());
+    }]>,
+  ];
 }
 
 // The 'LinalgStructuredInterface' provides access to the 'LinalgOp' interface.
 def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
   let cppNamespace = "::mlir::linalg";
   let methods = [
     //===------------------------------------------------------------------===//
     // Loop types handling.

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

 ods_def<MatmulI8I8I32Op>
 implements_interface<LinalgContractionOpInterface> :
 def matmul_i8_i8_i32(A: i8(M, K), B: i8(K, N)) -> (C: i32(M, N)) {
   // TODO: ideally something closer to
   //   C(m, n) += cast<i32>(A(m, k)) * cast<i32>(B(k, n))
   C(m, n) = std_addi<k>(std_sexti32(std_muli(A(m, k), B(k, n))));
 }
 
+ods_def<MatmulI16I16I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def matmul_i16_i16_i32(A: i16(M, K), B: i16(K, N)) -> (C: i32(M, N)) {
+  C(m, n) = std_addi<k>(std_sexti32(std_muli(A(m, k), B(k, n))));
+}
+
+ods_def<MatmulI32I32I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def matmul_i32_i32_i32(A: i32(M, K), B: i32(K, N)) -> (C: i32(M, N)) {
+  C(m, n) = std_addi<k>(std_muli(A(m, k), B(k, n)));
+}
+
 ods_def<MatvecOp>
 implements_interface<LinalgContractionOpInterface> :
 def matvec(A: f32(M, N), y: f32(N)) -> (x: f32(M)) {
   x(m) = std_addf<n>(std_mulf(A(m, n), y(n)));
 }
 
+ods_def<MatvecI8I8I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def matvec_i8_i8_i32(A: i8(M, N), y: i8(N)) -> (x: i32(M)) {
+  x(m) = std_addi<n>(std_sexti32(std_muli(A(m, n), y(n))));
+}
+
+ods_def<MatvecI16I16I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def matvec_i16_i16_i32(A: i16(M, N), y: i16(N)) -> (x: i32(M)) {
+  x(m) = std_addi<n>(std_sexti32(std_muli(A(m, n), y(n))));
+}
+
+ods_def<MatvecI32I32I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def matvec_i32_i32_i32(A: i32(M, N), y: i32(N)) -> (x: i32(M)) {
+  x(m) = std_addi<n>(std_muli(A(m, n), y(n)));
+}
+
 ods_def<VecmatOp>
 implements_interface<LinalgContractionOpInterface> :
 def vecmat(y: f32(M), A: f32(M, N)) -> (x: f32(N)) {
   x(n) = std_addf<m>(std_mulf(y(m), A(m, n)));
 }
 
+ods_def<VecmatI8I8I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def vecmat_i8_i8_i32(y: i8(M), A: i8(M, N)) -> (x: i32(N)) {
+  x(n) = std_addi<m>(std_sexti32(std_muli(y(m), A(m, n))));
+}
+
+ods_def<VecmatI16I16I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def vecmat_i16_i16_i32(y: i16(M), A: i16(M, N)) -> (x: i32(N)) {
+  x(n) = std_addi<m>(std_sexti32(std_muli(y(m), A(m, n))));
+}
+
+
+ods_def<VecmatI32I32I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def vecmat_i32_i32_i32(y: i32(M), A: i32(M, N)) -> (x: i32(N)) {
+  x(n) = std_addi<m>(std_muli(y(m), A(m, n)));
+}
+
 ods_def<DotOp>
 implements_interface<LinalgContractionOpInterface> :
 def dot(A: f32(M), B: f32(M)) -> (C: f32()) {
   C() = std_addf<m>(std_mulf(A(m), B(m)));
 }
 
+ods_def<DotI8I8I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def dot_i8_i8_i32(A: i8(M), B: i8(M)) -> (C: i32()) {
+  C() = std_addi<m>(std_sexti32(std_muli(A(m), B(m))));
+}
+
+ods_def<DotI16I16I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def dot_i16_i16_i32(A: i16(M), B: i16(M)) -> (C: i32()) {
+  C() = std_addi<m>(std_sexti32(std_muli(A(m), B(m))));
+}
+
+
+ods_def<DotI32I32I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def dot_i32_i32_i32(A: i32(M), B: i32(M)) -> (C: i32()) {
+  C() = std_addi<m>(std_muli(A(m), B(m)));
+}
+
+
 ods_def<BatchMatmulOp>
 implements_interface<LinalgContractionOpInterface> :
 def batch_matmul(A: f32(Batch, M, K), B: f32(Batch, K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(b, k, n)));
 }
 
+ods_def<BatchMatmulI8I8I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def batch_matmul_i8_i8_i32(A: i8(Batch, M, K), B: i8(Batch, K, N)) -> (C: i32(Batch, M, N)) {
+  C(b, m, n) = std_addi<k>(std_sexti32(std_muli(A(b, m, k), B(b, k, n))));
+}
+
+ods_def<BatchMatmulI16I16I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def batch_matmul_i16_i16_i32(A: i16(Batch, M, K), B: i16(Batch, K, N)) -> (C: i32(Batch, M, N)) {
+  C(b, m, n) = std_addi<k>(std_sexti32(std_muli(A(b, m, k), B(b, k, n))));
+}
+
+
+ods_def<BatchMatmulI32I32I32Op>
+implements_interface<LinalgContractionOpInterface> :
+def batch_matmul_i32_i32_i32(A: i32(Batch, M, K), B: i32(Batch, K, N)) -> (C: i32(Batch, M, N)) {
+  C(b, m, n) = std_addi<k>(std_muli(A(b, m, k), B(b, k, n)));
+}
+
 ods_def<ConvWOp>:
 def conv_1d(I: f32(W), K: f32(KW)) -> (O: f32(W)) {
   O(w) = std_addf<kw>(std_mulf(I(w + kw), K(kw)));
 }
 
 ods_def<ConvNWCOp>:
 def conv_1d_nwc(I: f32(N, W, C), K: f32(F, KW, C)) -> (O: f32(N, W, F)) {
   O(n, w, f) = std_addf<kw>(std_mulf(I(n, w + kw, c), K(f, kw, c)));