[PATCH] D80651: Harden MLIR detection of misconfiguration when missing dialect registration

[Mehdi AMINI via Phabricator via llvm-commits]

This revision was automatically updated to reflect the committed changes.
mehdi_amini marked an inline comment as done.
Closed by commit rG213c6cdf2e7a: Harden MLIR detection of misconfiguration when missing dialect registration (authored by mehdi_amini).

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D80651/new/

https://reviews.llvm.org/D80651

Files:
  mlir/include/mlir/IR/Builders.h
  mlir/include/mlir/IR/MLIRContext.h
  mlir/include/mlir/IR/OpDefinition.h
  mlir/lib/IR/MLIRContext.cpp


Index: mlir/lib/IR/MLIRContext.cpp
===================================================================
--- mlir/lib/IR/MLIRContext.cpp
+++ mlir/lib/IR/MLIRContext.cpp
@@ -543,6 +543,13 @@
   return result;
 }
 
+bool MLIRContext::isOperationRegistered(StringRef name) {
+  // Lock access to the context registry.
+  ScopedReaderLock registryLock(impl->contextMutex, impl->threadingIsEnabled);
+
+  return impl->registeredOperations.count(name);
+}
+
 void Dialect::addOperation(AbstractOperation opInfo) {
   assert((getNamespace().empty() ||
           opInfo.name.split('.').first == getNamespace()) &&
@@ -621,8 +628,9 @@
 static Dialect &lookupDialectForSymbol(MLIRContext *ctx, TypeID typeID) {
   auto &impl = ctx->getImpl();
   auto it = impl.registeredDialectSymbols.find(typeID);
-  assert(it != impl.registeredDialectSymbols.end() &&
-         "symbol is not registered.");
+  if (it == impl.registeredDialectSymbols.end())
+    llvm::report_fatal_error(
+        "Trying to create a type that was not registered in this MLIRContext.");
   return *it->second;
 }
 
Index: mlir/include/mlir/IR/OpDefinition.h
===================================================================
--- mlir/include/mlir/IR/OpDefinition.h
+++ mlir/include/mlir/IR/OpDefinition.h
@@ -1235,7 +1235,10 @@
   static bool classof(Operation *op) {
     if (auto *abstractOp = op->getAbstractOperation())
       return TypeID::get<ConcreteType>() == abstractOp->typeID;
-    return op->getName().getStringRef() == ConcreteType::getOperationName();
+    assert(op->getContext()->isOperationRegistered(
+               ConcreteType::getOperationName()) &&
+           "Casting attempt to an unregistered operation");
+    return false;
   }
 
   /// This is the hook used by the AsmParser to parse the custom form of this
Index: mlir/include/mlir/IR/MLIRContext.h
===================================================================
--- mlir/include/mlir/IR/MLIRContext.h
+++ mlir/include/mlir/IR/MLIRContext.h
@@ -85,6 +85,9 @@
   /// directly.
   std::vector<AbstractOperation *> getRegisteredOperations();
 
+  /// Return true if this operation name is registered in this context.
+  bool isOperationRegistered(StringRef name);
+
   // This is effectively private given that only MLIRContext.cpp can see the
   // MLIRContextImpl type.
   MLIRContextImpl &getImpl() { return *impl; }
Index: mlir/include/mlir/IR/Builders.h
===================================================================
--- mlir/include/mlir/IR/Builders.h
+++ mlir/include/mlir/IR/Builders.h
@@ -374,6 +374,10 @@
   template <typename OpTy, typename... Args>
   OpTy create(Location location, Args &&... args) {
     OperationState state(location, OpTy::getOperationName());
+    if (!state.name.getAbstractOperation())
+      llvm::report_fatal_error("Building op `" +
+                               state.name.getStringRef().str() +
+                               "` but it isn't registered in this MLIRContext");
     OpTy::build(*this, state, std::forward<Args>(args)...);
     auto *op = createOperation(state);
     auto result = dyn_cast<OpTy>(op);
@@ -390,6 +394,10 @@
     // Create the operation without using 'createOperation' as we don't want to
     // insert it yet.
     OperationState state(location, OpTy::getOperationName());
+    if (!state.name.getAbstractOperation())
+      llvm::report_fatal_error("Building op `" +
+                               state.name.getStringRef().str() +
+                               "` but it isn't registered in this MLIRContext");
     OpTy::build(*this, state, std::forward<Args>(args)...);
     Operation *op = Operation::create(state);
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D80651.266768.patch
Type: text/x-patch
Size: 3650 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200528/c949c81b/attachment.bin>