[Mlir-commits] [mlir] 4b897de - [mlir][ods] Add nested OpTrait

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2021-07-20T10:44:48-07:00
New Revision: 4b897de5fa0298398cb7e993963cc3507c961985

URL: https://github.com/llvm/llvm-project/commit/4b897de5fa0298398cb7e993963cc3507c961985
DIFF: https://github.com/llvm/llvm-project/commit/4b897de5fa0298398cb7e993963cc3507c961985.diff

LOG: [mlir][ods] Add nested OpTrait

Allows for grouping OpTraits with list of OpTrait to make it easier to group OpTraits together without needing to use list concats (e.g., enable using `[Traits, ..., UsefulGroupOfTraits, Others, ...]` instead of `[Traits, ...] # UsefulGroupOfTraits # [Others, ...]`). Flatten in construction of Operation. This recurses here as the expectation is that these aren't expected to be deeply nested (most likely only 1 level of nesting).

Differential Revision: https://reviews.llvm.org/D106223

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/lib/TableGen/Operator.cpp
    mlir/test/mlir-tblgen/op-decl-and-defs.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 59890190d2c66..1168dd12d6f09 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1796,6 +1796,14 @@ class PredTrait<string descr, Pred pred> : Trait {
 // TODO: Remove this class in favor of using Trait.
 class OpTrait;
 
+// Define a OpTrait corresponding to a list of OpTraits, this allows for
+// specifying a list of traits as trait. Avoids needing to do
+// `[Traits, ...] # ListOfTraits # [Others, ...]` while still allowing providing
+// convenient groupings.
+class OpTraitList<list<OpTrait> props> : OpTrait {
+  list<OpTrait> traits = props;
+}
+
 // These classes are used to define operation specific traits.
 class NativeOpTrait<string name> : NativeTrait<name, "Op">, OpTrait;
 class ParamNativeOpTrait<string prop, string params>

diff  --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index 11f95c0957ef8..ea9513d4e6647 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -489,11 +489,21 @@ void Operator::populateOpStructure() {
     // This is uniquing based on pointers of the trait.
     SmallPtrSet<const llvm::Init *, 32> traitSet;
     traits.reserve(traitSet.size());
-    for (auto *traitInit : *traitList) {
-      // Keep traits in the same order while skipping over duplicates.
-      if (traitSet.insert(traitInit).second)
-        traits.push_back(Trait::create(traitInit));
-    }
+
+    std::function<void(llvm::ListInit *)> insert;
+    insert = [&](llvm::ListInit *traitList) {
+      for (auto *traitInit : *traitList) {
+        auto *def = cast<DefInit>(traitInit)->getDef();
+        if (def->isSubClassOf("OpTraitList")) {
+          insert(def->getValueAsListInit("traits"));
+          continue;
+        }
+        // Keep traits in the same order while skipping over duplicates.
+        if (traitSet.insert(traitInit).second)
+          traits.push_back(Trait::create(traitInit));
+      }
+    };
+    insert(traitList);
   }
 
   populateTypeInferenceInfo(argumentsAndResultsIndex);

diff  --git a/mlir/test/mlir-tblgen/op-decl-and-defs.td b/mlir/test/mlir-tblgen/op-decl-and-defs.td
index 4fb9ecb39730a..471bac6edf858 100644
--- a/mlir/test/mlir-tblgen/op-decl-and-defs.td
+++ b/mlir/test/mlir-tblgen/op-decl-and-defs.td
@@ -261,6 +261,15 @@ def NS_JOp : NS_Op<"op_with_InferTypeOpInterface_interface", [DeclareOpInterface
 // CHECK: static void build(::mlir::OpBuilder &, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::ValueRange operands, ::llvm::ArrayRef<::mlir::NamedAttribute> attributes = {});
 
+// Test usage of OpTraitList getting flattened during emission.
+def NS_KOp : NS_Op<"k_op", [IsolatedFromAbove,
+    OpTraitList<[DeclareOpInterfaceMethods<InferTypeOpInterface>]>]> {
+  let arguments = (ins AnyType:$a, AnyType:$b);
+}
+
+// CHECK: class KOp : public ::mlir::Op<KOp,
+// CHECK-SAME: ::mlir::OpTrait::IsIsolatedFromAbove, ::mlir::InferTypeOpInterface::Trait
+
 // Check native OpTrait usage
 // ---