Add contract to mlir-gen (#1027)

Adds `linalg.contract` to MLIR Gen so we can start testing contraction
inputs.

Author: rengolin

test/Integration/mlir-gen.mlir

@@ -8,10 +8,18 @@
 // RUN: mlir-gen --kernel=const --bias --relu --seed=123 --batch=10 --layers=10,10,10 | tpp-run -e entry -entry-point-result=void
 
 // Matmul only
-// RUN: mlir-gen --kernel=const --bias --relu --seed=123 --batch=10 --layers=10,10 | tpp-run -e entry -entry-point-result=void
+// RUN: mlir-gen --kernel=const --batch=10 --layers=10,10 | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=MATMUL
+// RUN: mlir-gen --kernel=const --batch=10 --layers=10,10 --output=generic | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=MATMUL
+// RUN: mlir-gen --kernel=const --batch=10 --layers=10,10 --output=contract | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=MATMUL
+// RUN: mlir-gen --kernel=const --batch=10 --layers=10,10 --output=named | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=MATMUL
+// RUN: mlir-gen --kernel=const --batch=10 --layers=10,10 --output=named --keep-generic-matmul | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=MATMUL
 
 // Constant values
 // RUN: mlir-gen --kernel=const --bias --relu --batch=10 --layers=10,10 | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=CONSTANT
+// RUN: mlir-gen --kernel=const --bias --relu --batch=10 --layers=10,10 --output=generic | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=CONSTANT
+// RUN: mlir-gen --kernel=const --bias --relu --batch=10 --layers=10,10 --output=contract | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=CONSTANT
+// RUN: mlir-gen --kernel=const --bias --relu --batch=10 --layers=10,10 --output=named | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=CONSTANT
+// RUN: mlir-gen --kernel=const --bias --relu --batch=10 --layers=10,10 --output=named --keep-generic-matmul | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=CONSTANT
 
 // Kernel - matmul
 // RUN: mlir-gen --kernel=args --seed=123 --float-type=f32 --batch=10 --layers=10,10 | tpp-run -e entry -entry-point-result=void -print | FileCheck %s --check-prefix=GEN-MATMUL
@@ -23,6 +31,8 @@
 // RUN: mlir-gen --kernel=const --bias --relu --seed=123 --batch=10 --layers=10,10 --tiles=2,2,2 | tpp-run -e entry -entry-point-result=void -n 10 | FileCheck %s --check-prefix=PERF
 // RUN: mlir-gen --kernel=const --bias --relu --seed=123 --batch=10 --layers=10,10,10 --tiles=2,2,2 | tpp-run -e entry -entry-point-result=void -n 10 | FileCheck %s --check-prefix=PERF
 
+// MATMUL:( 10, 10, 10, 10, 10, 10, 10, 10, 10, 10 )
+
 // CONSTANT:( 11, 11, 11, 11, 11, 11, 11, 11, 11, 11 )
 
 // GEN-MATMUL: ( 11, 11, 11, 11, 11, 11, 11, 11, 11, 11 )

tools/mlir-gen/MLIRGen.cpp

@@ -85,9 +85,12 @@ MLIRGenerator::MLIRGenerator(StringRef outputOpKindStr, StringRef kernelStr,
   auto optOutputOpKind =
       llvm::StringSwitch<std::optional<OutputOpKind>>(outputOpKindStr)
           .CaseLower("generic", OutputOpKind::Generic)
+          .CaseLower("contract", OutputOpKind::Contract)
           .CaseLower("named", OutputOpKind::NamedOp)
           .Default(std::nullopt);
   assert(optOutputOpKind && "Invalid output Op kind");
+  assert(!(optOutputOpKind == OutputOpKind::Contract && keepGenericMatmul) &&
+         "Can't keep generic matmul with contract");
   outputOpKind = *optOutputOpKind;
 
   // Parse kernel type
@@ -181,7 +184,7 @@ Value MLIRGenerator::createLayer(LayerArgs &args) {
   if (outputOpKind == OutputOpKind::Generic) {
     chain = lowerBiasAdd(chain, args.bias.value, args.output.value);
     chain = lowerRelu(chain, args.output.value);
-  } else if (outputOpKind == OutputOpKind::NamedOp) {
+  } else {
     chain = lowerNamedBiasAdd(chain, args.bias.value, args.output.value);
     chain = lowerNamedRelu(chain, args.output.value);
   }
@@ -190,7 +193,7 @@ Value MLIRGenerator::createLayer(LayerArgs &args) {
   if (args.index == layers.size() - 1) {
     if (outputOpKind == OutputOpKind::Generic) {
       chain = lowerSoftmax(chain, args.output.value);
-    } else if (outputOpKind == OutputOpKind::NamedOp) {
+    } else {
       chain = lowerNamedSoftmax(chain, args.output.value);
     }
   }
@@ -405,9 +408,10 @@ Value MLIRGenerator::lowerMatmul(Value input, Value weight, Value output) {
                                                   reassociationIndices);
   }
 
-  if (outputOpKind == OutputOpKind::Generic ||
-      (outputOpKind == OutputOpKind::NamedOp && keepGenericMatmul)) {
+  if (outputOpKind == OutputOpKind::Generic || keepGenericMatmul) {
     chain = lowerGenericMatmul(input, weight, output);
+  } else if (outputOpKind == OutputOpKind::Contract) {
+    chain = lowerContract(input, weight, output);
   } else if (outputOpKind == OutputOpKind::NamedOp) {
     chain = lowerNamedMatmul(input, weight, output);
   }
@@ -442,6 +446,21 @@ Value MLIRGenerator::lowerGenericMatmul(Value input, Value weight,
   return matmul;
 }
 
+Value MLIRGenerator::lowerContract(Value input, Value weight, Value output) {
+  // Matmul as a linalg.contract
+  SmallVector<Attribute> maps;
+  maps.push_back(AffineMapAttr::get(getMap(input, MAP_MATMUL_INPUT)));   // { 0, 2 }
+  maps.push_back(AffineMapAttr::get(getMap(weight, MAP_MATMUL_WEIGHT))); // { 2, 1 }
+  maps.push_back(AffineMapAttr::get(getMap(output, MAP_MATMUL_OUTPUT))); // { 0, 1 }
+  auto contract = builder
+                      .create<linalg::ContractOp>(
+                          loc, output.getType(), ValueRange{input, weight}, ValueRange{output},
+                          builder.getArrayAttr(maps))
+                      .getResult(0);
+
+  return contract;
+}
+
 Value MLIRGenerator::lowerBiasAdd(Value input, Value bias, Value output) {
   if (!enableBias)
     return input;

tools/mlir-gen/MLIRGen.h

@@ -75,7 +75,7 @@ class MLIRGenerator {
   bool enableSoftmax;
 
   /// List of linalg output Op kind which can be generated
-  enum class OutputOpKind { Generic, NamedOp };
+  enum class OutputOpKind { Generic, Contract, NamedOp };
 
   /// Kind of linalg output Op to be generated
   OutputOpKind outputOpKind;
@@ -156,6 +156,9 @@ class MLIRGenerator {
   /// Creates linalg named matmul
   Value lowerNamedMatmul(Value, Value, Value);
 
+  /// Creates linalg contract
+  Value lowerContract(Value, Value, Value);
+
   /// Creates a bias add in the current function
   /// Args: Input, Output (same for in-place)
   /// Returns the chain value to be used in the next op

tools/mlir-gen/mlir-gen.cpp

@@ -33,8 +33,8 @@ using namespace mlir;
 
 // Kind of linalg Op, generic or nameed ops
 llvm::cl::opt<std::string> outputOpKind(
-    "output", llvm::cl::desc("Specifies linalg op kind generic or named"),
-    llvm::cl::value_desc("generic,named"), llvm::cl::init("generic"));
+    "output", llvm::cl::desc("Specifies linalg op kind generic, contract or named"),
+    llvm::cl::value_desc("generic,contract,named"), llvm::cl::init("generic"));
 
 // Enable emission of generic matmul when outputKind is named op
 llvm::cl::opt<bool> keepGenericMatmul(