[mlir][ods] Add nested OpTrait

[jpienaar]

Summary

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).

[rriddle]

mlir/include/mlir/IR/OpBase.td

1802

What about naming this something like OpTraitList?

[jpienaar]

Updated to use OpTraitList

[jpienaar]

mlir/include/mlir/IR/OpBase.td

1802

Good idea

[Chia-hungDuan]

mlir/include/mlir/IR/OpBase.td

1800

typo?

mlir/test/mlir-tblgen/op-decl-and-defs.td

264–265

Does this need to be updated? Or the case below needs one more IsolatedFromAbove in the OpTraitList?

[jpienaar]

Fix typo and comment documenting test case

mlir/include/mlir/IR/OpBase.td

 //===----------------------------------------------------------------------===//
 // OpTrait definitions
 //===----------------------------------------------------------------------===//
 
 // OpTrait represents a trait regarding an operation.
 // 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>
     : ParamNativeTrait<prop, params, "Op">, OpTrait;
 class GenInternalOpTrait<string prop> : GenInternalTrait<prop, "Op">, OpTrait;
 class PredOpTrait<string descr, Pred pred> : PredTrait<descr, pred>, OpTrait;
 
 // Op defines an affine scope.

mlir/lib/TableGen/Operator.cpp

void Operator::populateOpStructure() {
   }
 
   // Create list of traits, skipping over duplicates: appending to lists in
   // tablegen is easy, making them unique less so, so dedupe here.
   if (auto *traitList = def.getValueAsListInit("traits")) {
     // This is uniquing based on pointers of the trait.
     SmallPtrSet<const llvm::Init *, 32> traitSet;
     traits.reserve(traitSet.size());
+
+    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);
 
   // Handle regions
   auto *regionsDag = def.getValueAsDag("regions");
   auto *regionsOp = dyn_cast<DefInit>(regionsDag->getOperator());
   if (!regionsOp || regionsOp->getDef()->getName() != "region") {

mlir/test/mlir-tblgen/op-decl-and-defs.td

 }
 // CHECK_LABEL: class NS_JOp :
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Type r, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::Value a, ::mlir::Value b);
 // CHECK: static void build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange resultTypes, ::mlir::Value a, ::mlir::Value b);
 // 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
 // ---
 
 def NS_TestTrait : NativeOpTrait<"TestTrait"> {
     let cppNamespace = "SomeNamespace";
 }
 
 def NS_KWithTraitOp : NS_Op<"KWithTrait", [NS_TestTrait]>;