From 5867c18e2c0d403b51a594897d56d935286748e4 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Fri, 26 Feb 2021 13:01:03 -0800
Subject: [PATCH] [mlir][linalg] Generate additional interfaces for named ops.

* Adds ContractionOpInterface to polymorphic_matmul.

Differential Revision: https://reviews.llvm.org/D97601
---
 .../mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml       |  2 ++
 mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp    | 10 +++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 02f0814..43e41ab 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -7,6 +7,8 @@ metadata: !LinalgOpMetadata
 
     This op is presently here to test a new path for generation and will replace
     the existing 'matmul' op when ready. Do not use.
+  implements:
+  - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !<LinalgTensorDef>
diff --git a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
index 5925e55..f651b03 100644
--- a/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
+++ b/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
@@ -53,6 +53,7 @@ struct LinalgOpMetadata {
   std::string name;
   std::string cppOpName;
   Optional<std::string> doc;
+  SmallVector<std::string> implements;
 };
 
 struct SerializedAffineMap {
@@ -197,6 +198,7 @@ struct MappingTraits<LinalgOpMetadata> {
     io.mapRequired("name", info.name);
     io.mapRequired("cpp_op_name", info.cppOpName);
     io.mapOptional("doc", info.doc);
+    io.mapOptional("implements", info.implements);
   }
 };
 
@@ -387,11 +389,11 @@ static const char structuredOpOdsHeaderFormat[] = R"FMT(
 // Op definition for {0}
 //===----------------------------------------------------------------------===//
 
-def {0} : LinalgStructuredBase_Op<"{1}", [
+def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([
   AttrSizedOperandSegments,
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-  SingleBlockImplicitTerminator<"YieldOp">
-  /*extraInterfaces=*/{2}]> {
+  SingleBlockImplicitTerminator<"YieldOp">],
+  /*extraInterfaces=*/[{2}])> {
     {3}
     let arguments = (ins
       Variadic<AnyShaped>:$inputs,
@@ -528,6 +530,8 @@ static LogicalResult generateNamedGenericOpOds(LinalgOpConfig &opConfig,
     doc = llvm::formatv(docFmt, summary.trim(), description.trim());
   }
 
+  interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");
+
   os << llvm::formatv(structuredOpOdsHeaderFormat, opConfig.metadata->cppOpName,
                       opConfig.metadata->name, interfaceNameList, doc, attrList,
                       opConfig.structuredOp->args.size(), attrBuilder,
-- 
2.7.4