[mlir][ODS] Allow to specify custom namespace for `NativeOpTrait`

[vinograd47]

Summary

This will allow to use NativeOpTrait and Operations
declared outside of mlir namespace.

mlir/include/mlir/IR/OpBase.td

 //===----------------------------------------------------------------------===//
 
 // OpTrait represents a trait regarding an op.
 class OpTrait;
 
 // NativeOpTrait corresponds to the MLIR C++ OpTrait mechanism. The
 // purpose to wrap around C++ symbol string with this class is to make
 // traits specified for ops in TableGen less alien and more integrated.
-class NativeOpTrait<string prop> : OpTrait {
-  string trait = "::mlir::OpTrait::" # prop;
+class NativeOpTrait<string name> : OpTrait {
+  string trait = name;
+  string cppNamespace = "::mlir::OpTrait";
 }
 
 // ParamNativeOpTrait corresponds to the template-parameterized traits in the
 // C++ implementation.  MLIR uses nested class templates to implement such
 // traits leading to constructs of the form "TraitName<Parameters>::Impl". Use
 // the value in `prop` as the trait name and the value in `params` as
 // parameters to construct the native trait class name.
 class ParamNativeOpTrait<string prop, string params>

 }
 
 // OpInterfaceTrait corresponds to a specific 'OpInterface' class defined in
 // C++. The purpose to wrap around C++ symbol string with this class is to make
 // interfaces specified for ops in TableGen less alien and more integrated.
 class OpInterfaceTrait<string name, code verifyBody = [{}]>
           : NativeOpTrait<""> {
   let trait = name # "::Trait";
+  let cppNamespace = "";
 
   // Specify the body of the verification function. `$_op` will be replaced with
   // the operation being verified.
   code verify = verifyBody;
 
   // An optional code block containing extra declarations to place in the
   // interface trait declaration.
   code extraTraitClassDeclaration = "";

mlir/include/mlir/TableGen/OpTrait.h

protected:
   const llvm::Record *def;
   Kind kind;
 };
 
 // OpTrait corresponding to a native C++ OpTrait.
 class NativeOpTrait : public OpTrait {
 public:
   // Returns the trait corresponding to a C++ trait class.
-  StringRef getTrait() const;
+  std::string getTrait() const;
 
   static bool classof(const OpTrait *t) { return t->getKind() == Kind::Native; }
 };
 
 // OpTrait corresponding to a predicate on the operation.
 class PredOpTrait : public OpTrait {
 public:
   // Returns the template for constructing the predicate.

mlir/lib/TableGen/OpTrait.cpp

OpTrait OpTrait::create(const llvm::Init *init) {
   if (def->isSubClassOf("OpInterfaceTrait"))
     return OpTrait(Kind::Interface, def);
   assert(def->isSubClassOf("NativeOpTrait"));
   return OpTrait(Kind::Native, def);
 }
 
 OpTrait::OpTrait(Kind kind, const llvm::Record *def) : def(def), kind(kind) {}
 
-llvm::StringRef NativeOpTrait::getTrait() const {
-  return def->getValueAsString("trait");
+std::string NativeOpTrait::getTrait() const {
+  llvm::StringRef trait = def->getValueAsString("trait");
+  llvm::StringRef cppNamespace = def->getValueAsString("cppNamespace");
+  return cppNamespace.empty() ? trait.str()
+                              : (cppNamespace + "::" + trait).str();
 }
 
 llvm::StringRef InternalOpTrait::getTrait() const {
   return def->getValueAsString("trait");
 }
 
 std::string PredOpTrait::getPredTemplate() const {
   auto pred = Pred(def->getValueInit("predicate"));

mlir/test/mlir-tblgen/op-decl.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 = {});
 
+// Check native OpTrait usage
+// ---
+
+def NS_TestTrait : NativeOpTrait<"TestTrait"> {
+    let cppNamespace = "SomeNamespace";
+}
+
+def NS_KWithTraitOp : NS_Op<"KWithTrait", [NS_TestTrait]>;
+
+// CHECK-LABEL: NS::KWithTraitOp declarations
+// CHECK:       class KWithTraitOp : public ::mlir::Op<KWithTraitOp
+// CHECK-SAME:      SomeNamespace::TestTrait
+
 // Test that type defs have the proper namespaces when used as a constraint.
 // ---
 
 def Test_Dialect2 : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::dialect2";
 }
 def TestDialect2Type : TypeDef<Test_Dialect2, "Dialect2Type">;