Harden MLIR detection of misconfiguration when missing dialect registration

[rriddle]

I wonder if we should also add something to Op::classof. i.e. can you cast to FooOp if it isn't registered?

[mehdi_amini]

In D80651#2058483, @rriddle wrote:

I wonder if we should also add something to Op::classof. i.e. can you cast to FooOp if it isn't registered?

Right now I was trying to add code that runs in release mode, so I tried to hook into code path that are OK to pay the price: type creation is once per type and op creation is expensive enough to take a nullptr check.
Are you proposing to have an extra nullptr check for every cast to an operation? I can also do that.

[jpienaar]

In D80651#2058607, @mehdi_amini wrote:

In D80651#2058483, @rriddle wrote:

I wonder if we should also add something to Op::classof. i.e. can you cast to FooOp if it isn't registered?

Right now I was trying to add code that runs in release mode, so I tried to hook into code path that are OK to pay the price: type creation is once per type and op creation is expensive enough to take a nullptr check.
Are you proposing to have an extra nullptr check for every cast to an operation? I can also do that.

I'd prefer if we do keep it light: it would seem that if one is testing this with unit tests/NDEBUG then these would be caught with asserts . So with proper unit tests one should be catching these & should be rare (e.g., how often is a new dialect added to production build?) [which is also why I'm a little bit torn but don't feel strongly against].

[rriddle]

In D80651#2058617, @jpienaar wrote:

In D80651#2058607, @mehdi_amini wrote:

In D80651#2058483, @rriddle wrote:

I wonder if we should also add something to Op::classof. i.e. can you cast to FooOp if it isn't registered?

Right now I was trying to add code that runs in release mode, so I tried to hook into code path that are OK to pay the price: type creation is once per type and op creation is expensive enough to take a nullptr check.
Are you proposing to have an extra nullptr check for every cast to an operation? I can also do that.

I'd prefer if we do keep it light: it would seem that if one is testing this with unit tests/NDEBUG then these would be caught with asserts . So with proper unit tests one should be catching these & should be rare (e.g., how often is a new dialect added to production build?) [which is also why I'm a little bit torn but don't feel strongly against].

I brought this up because if we have a policy of not allowing the use of *Op classes for unregistered ops we should be consistent with erroring out.

As for enabling in release builds it depends on how often we expect unregistered operations to be used in production. If it is often, then we should only enable during debug builds. If not, it is fairly cheap to enable during release given that classof already has a fast path for operations that are registered. Regardless of which one we choose, it has the benefit of removing the fallback string comparison in favor of always returning false.

[mehdi_amini]

Disable cast<OpTy>(op) on unregistered operations

[mehdi_amini]

Add unreachable for the cast<OpTy>

[mehdi_amini]

(build fix)

[mehdi_amini]

I thought it would involve adding an extra check, but if it is free on the fast path then checking the cast<> case seems good!

We had many issues where the unit-test environment does not reproduce the exact registration state of the production path, this will help saving a bunch of debugging hopefully

[rriddle]

mlir/include/mlir/IR/OpDefinition.h

1238–1241

You will still want to check here. This path is only taken if the operation being casted is unregistered, not ConcreteType itself.

[mehdi_amini]

Add a method on the MLIRContext to check if an operation name is registered in the context

[mehdi_amini]

mlir/include/mlir/IR/OpDefinition.h

1238–1241

Yeah I figured when the test were failing...
So unfortunately this isn't "free" and I had to make it an assertion.

mlir/include/mlir/IR/Builders.h

public:
 
   /// Creates an operation given the fields represented as an OperationState.
   Operation *createOperation(const OperationState &state);
 
   /// Create an operation of specific op type at the current insertion point.
   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);
     assert(result && "builder didn't return the right type");
     return result;
   }
 
   /// Create an operation of specific op type at the current insertion point,
   /// and immediately try to fold it. This functions populates 'results' with
   /// the results after folding the operation.
   template <typename OpTy, typename... Args>
   void createOrFold(SmallVectorImpl<Value> &results, Location location,
                     Args &&... args) {
     // 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);
 
     // Fold the operation. If successful destroy it, otherwise insert it.
     if (succeeded(tryFold(op, results)))
       op->destroy();
     else
       insert(op);

mlir/include/mlir/IR/MLIRContext.h

public:
   /// emitting diagnostics.
   void printStackTraceOnDiagnostic(bool enable);
 
   /// Return information about all registered operations.  This isn't very
   /// efficient: typically you should ask the operations about their properties
   /// 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; }
 
   /// Returns the diagnostic engine for this context.
   DiagnosticEngine &getDiagEngine();
 
   /// Returns the storage uniquer used for creating affine constructs.

mlir/include/mlir/IR/OpDefinition.h

public:
 
   /// Return the parent Region of this operation.
   Region *getParentRegion() { return getOperation()->getParentRegion(); }
 
   /// Return true if this "op class" can match against the specified operation.
   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
   /// op from an .mlir file.  Op implementations should provide a parse method,
   /// which returns failure.  On success, they should return fill in result with
   /// the fields to use.
   static ParseResult parseAssembly(OpAsmParser &parser,
                                    OperationState &result) {

mlir/lib/IR/MLIRContext.cpp

std::vector<AbstractOperation *> MLIRContext::getRegisteredOperations() {
 
   std::vector<AbstractOperation *> result;
   result.reserve(opsToSort.size());
   for (auto &elt : opsToSort)
     result.push_back(elt.second);
   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()) &&
          "op name doesn't start with dialect namespace");
   assert(&opInfo.dialect == this && "Dialect object mismatch");
   auto &impl = context->getImpl();
 
   // Lock access to the context registry.

 
 //===----------------------------------------------------------------------===//
 // Type uniquing
 //===----------------------------------------------------------------------===//
 
 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;
 }
 
 /// Returns the storage uniquer used for constructing type storage instances.
 /// This should not be used directly.
 StorageUniquer &MLIRContext::getTypeUniquer() { return getImpl().typeUniquer; }
 
 /// Get the dialect that registered the type with the provided typeid.