[Dialect] [Linalgx] Add linalgx ops: 3 vnni matmuls and multi_batch_matmul (#89)

Author: Longsheng Du

include/gc/Dialect/Linalgx/LinalgxStructuredOps.td

@@ -104,4 +104,212 @@ def Linalgx_SigmoidOp : LinalgxStructuredBase_Op<"sigmoid",
   }];
 }
 
+def Linalgx_Mm2DVnniOp 
+    : LinalgxStructuredBase_Op<"mm2d_vnni", [AttrSizedOperandSegments]> {
+  let summary = "Transposed matmul with 2d input and vnni packed weights";
+  let description = [{
+    Supported format: A[M, K] * B[N0, K0, k, n, v] -> C[M, N], with:
+      N = N0 * n
+      K = K0 * k * v; v = (2, 4)
+  }];
+  let arguments = (ins
+      Variadic<TensorOrMemref>:$inputs,
+      Variadic<TensorOrMemref>:$outputs);
+  let results = (outs Variadic<TensorOrMemref>:$results);
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+      OpBuilder<
+      (ins 
+        "TypeRange":$resultTensorTypes, 
+        "ValueRange":$inputs,
+        "ValueRange":$outputs,
+        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildStructuredOp($_builder, $_state, resultTensorTypes,
+          inputs, outputs, attributes, Mm2DVnniOp::getRegionBuilder());
+      }]>
+  ];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    // Declare functions necessary for LinalgStructuredInterface.
+    SmallVector<utils::IteratorType> getIteratorTypesArray();
+    ArrayAttr getIndexingMaps();
+    static unsigned getNumRegionArgs() { return 3; }
+    std::string getLibraryCallName() {
+      return "op_has_no_registered_library_name";
+    }
+
+    // Implement functions necessary for DestinationStyleOpInterface.
+    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
+
+    static void regionBuilder(ImplicitLocOpBuilder &b,
+                              Block &block, ArrayRef<NamedAttribute> attrs);
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() {
+      return regionBuilder;
+    }
+  }];
+}
+
+def Linalgx_Mm4DVnniOp 
+    : LinalgxStructuredBase_Op<"mm4d_vnni", [AttrSizedOperandSegments]> {
+  let summary = "Transposed matmul with 4d blocking input and vnni packed weights";
+  let description = [{
+    Supported format: A[M, K, m, k] * B[N, K, k0, n, v] -> C[M, N, m, n], with:
+      k = k0 * v; v = (2, 4)
+  }];
+  let arguments = (ins
+      Variadic<TensorOrMemref>:$inputs,
+      Variadic<TensorOrMemref>:$outputs);
+  let results = (outs Variadic<TensorOrMemref>:$results);
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+      OpBuilder<
+      (ins 
+        "TypeRange":$resultTensorTypes, 
+        "ValueRange":$inputs,
+        "ValueRange":$outputs,
+        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildStructuredOp($_builder, $_state, resultTensorTypes,
+          inputs, outputs, attributes, Mm4DVnniOp::getRegionBuilder());
+      }]>
+  ];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    // Declare functions necessary for LinalgStructuredInterface.
+    SmallVector<utils::IteratorType> getIteratorTypesArray();
+    ArrayAttr getIndexingMaps();
+    static unsigned getNumRegionArgs() { return 3; }
+    std::string getLibraryCallName() {
+      return "op_has_no_registered_library_name";
+    }
+
+    // Implement functions necessary for DestinationStyleOpInterface.
+    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
+
+    static void regionBuilder(ImplicitLocOpBuilder &b,
+                              Block &block, ArrayRef<NamedAttribute> attrs);
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() {
+      return regionBuilder;
+    }
+  }];
+}
+
+def Linalgx_BatchReduceMatmulVnniOp 
+    : LinalgxStructuredBase_Op<"batch_reduce_matmul_vnni", [AttrSizedOperandSegments]> {
+  let summary = "Batch reduced matmul with 3d batch input and vnni packed weights";
+  let description = [{
+    Supported format: A[B, M, K] * B[B, k, N, v] -> C[M, N], with:
+      K = k * v; v = (2, 4)
+  }];
+  let arguments = (ins
+      Variadic<TensorOrMemref>:$inputs,
+      Variadic<TensorOrMemref>:$outputs);
+  let results = (outs Variadic<TensorOrMemref>:$results);
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+      OpBuilder<
+      (ins 
+        "TypeRange":$resultTensorTypes, 
+        "ValueRange":$inputs,
+        "ValueRange":$outputs,
+        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildStructuredOp($_builder, $_state, resultTensorTypes,
+          inputs, outputs, attributes, BatchReduceMatmulVnniOp::getRegionBuilder());
+      }]>
+  ];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    // Declare functions necessary for LinalgStructuredInterface.
+    SmallVector<utils::IteratorType> getIteratorTypesArray();
+    ArrayAttr getIndexingMaps();
+    static unsigned getNumRegionArgs() { return 3; }
+    std::string getLibraryCallName() {
+      return "op_has_no_registered_library_name";
+    }
+
+    // Implement functions necessary for DestinationStyleOpInterface.
+    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
+
+    static void regionBuilder(ImplicitLocOpBuilder &b,
+                              Block &block, ArrayRef<NamedAttribute> attrs);
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() {
+      return regionBuilder;
+    }
+  }];
+}
+
+def Linalgx_MultiBatchMatmulOp : LinalgxStructuredBase_Op<"multi_batch_matmul", 
+      [AttrSizedOperandSegments, LinalgContractionOpInterface]> {
+  let summary = "Batch matmul with variable batch dims";
+  let arguments = (ins
+      Variadic<TensorOrMemref>:$inputs,
+      Variadic<TensorOrMemref>:$outputs);
+  let results = (outs Variadic<TensorOrMemref>:$results);
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+      OpBuilder<
+      (ins 
+        "TypeRange":$resultTensorTypes, 
+        "ValueRange":$inputs,
+        "ValueRange":$outputs,
+        CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildStructuredOp($_builder, $_state, resultTensorTypes,
+          inputs, outputs, attributes, MultiBatchMatmulOp::getRegionBuilder());
+      }]>
+  ];
+
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    // Declare functions necessary for LinalgStructuredInterface.
+    SmallVector<utils::IteratorType> getIteratorTypesArray();
+    ArrayAttr getIndexingMaps();
+    static unsigned getNumRegionArgs() { return 3; }
+    std::string getLibraryCallName() {
+      return "op_has_no_registered_library_name";
+    }
+
+    // Implement functions necessary for DestinationStyleOpInterface.
+    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
+
+    static void regionBuilder(ImplicitLocOpBuilder &b,
+                              Block &block, ArrayRef<NamedAttribute> attrs);
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() {
+      return regionBuilder;
+    }
+  }];
+}
+
 #endif // LINALGX_STRUCTURED_OPS