[Mlir-commits] [mlir] 035abe3 - [mlir][ODS] Allow to specify custom namespace for `NativeOpTrait`

[Vladislav Vinogradov]

Author: Vladislav Vinogradov
Date: 2021-02-08T10:25:45+03:00
New Revision: 035abe30c9a55c08229f53aa63745a9ba87892d1

URL: https://github.com/llvm/llvm-project/commit/035abe30c9a55c08229f53aa63745a9ba87892d1
DIFF: https://github.com/llvm/llvm-project/commit/035abe30c9a55c08229f53aa63745a9ba87892d1.diff

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

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

Reviewed By: ftynse

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/include/mlir/TableGen/OpTrait.h
    mlir/lib/TableGen/OpTrait.cpp
    mlir/test/mlir-tblgen/op-decl.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 0bd1ee38bf1f..31598c143211 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1718,8 +1718,9 @@ 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
@@ -1852,6 +1853,7 @@ class CArg<string ty, string value = ""> {
 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.

diff  --git a/mlir/include/mlir/TableGen/OpTrait.h b/mlir/include/mlir/TableGen/OpTrait.h
index 8af7b6af1017..4ac0c1c4ed25 100644
--- a/mlir/include/mlir/TableGen/OpTrait.h
+++ b/mlir/include/mlir/TableGen/OpTrait.h
@@ -63,7 +63,7 @@ class 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; }
 };

diff  --git a/mlir/lib/TableGen/OpTrait.cpp b/mlir/lib/TableGen/OpTrait.cpp
index 793596c98266..4ebfc9bb531b 100644
--- a/mlir/lib/TableGen/OpTrait.cpp
+++ b/mlir/lib/TableGen/OpTrait.cpp
@@ -35,8 +35,11 @@ OpTrait OpTrait::create(const llvm::Init *init) {
 
 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 {

diff  --git a/mlir/test/mlir-tblgen/op-decl.td b/mlir/test/mlir-tblgen/op-decl.td
index 91e46b59c2e6..f15c5c6a6a1b 100644
--- a/mlir/test/mlir-tblgen/op-decl.td
+++ b/mlir/test/mlir-tblgen/op-decl.td
@@ -248,6 +248,19 @@ 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 = {});
 
+// 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.
 // ---