[mlir] Add C and python API for is_registered_operation.

[stellaraccident]

Summary

• Suggested to be broken out of D99578

[stellaraccident]

Backout weird clang format changes.

[stellaraccident]

Backout more clang-format stuff.

[mehdi_amini]

mlir/include/mlir-c/IR.h

124

[stellaraccident]

Comments.

[ftynse]

mlir/test/Bindings/Python/dialects.py

33

Nice!

mlir/include/mlir-c/IR.h

 /// Gets the dialect instance owned by the given context using the dialect
 /// namespace to identify it, loads (i.e., constructs the instance of) the
 /// dialect if necessary. If the dialect is not registered with the context,
 /// returns null. Use mlirContextLoad<Name>Dialect to load an unregistered
 /// dialect.
 MLIR_CAPI_EXPORTED MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                                            MlirStringRef name);
 
+/// Returns whether the given fully-qualified operation (i.e.
+/// 'dialect.operation') is registered with the context. This will return true
+/// if the dialect is registered and the operation is registered within the
+/// dialect.
+MLIR_CAPI_EXPORTED bool mlirContextIsRegisteredOperation(MlirContext context,
+                                                         MlirStringRef name);
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
 
 /// Returns the context that owns the dialect.
 MLIR_CAPI_EXPORTED MlirContext mlirDialectGetContext(MlirDialect dialect);
 
 /// Checks if the dialect is null.

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyMlirContext>(m, "Context")
           "Gets or loads a dialect by name, returning its descriptor object")
       .def_property(
           "allow_unregistered_dialects",
           [](PyMlirContext &self) -> bool {
             return mlirContextGetAllowUnregisteredDialects(self.get());
           },
           [](PyMlirContext &self, bool value) {
             mlirContextSetAllowUnregisteredDialects(self.get(), value);
+          })
+      .def("is_registered_operation",
+           [](PyMlirContext &self, std::string &name) {
+             return mlirContextIsRegisteredOperation(
+                 self.get(), MlirStringRef{name.data(), name.size()});
            });
 
   //----------------------------------------------------------------------------
   // Mapping of PyDialectDescriptor
   //----------------------------------------------------------------------------
   py::class_<PyDialectDescriptor>(m, "DialectDescriptor")
       .def_property_readonly("namespace",
                              [](PyDialectDescriptor &self) {
                                MlirStringRef ns =

mlir/lib/CAPI/IR/IR.cpp

intptr_t mlirContextGetNumLoadedDialects(MlirContext context) {
   return static_cast<intptr_t>(unwrap(context)->getLoadedDialects().size());
 }
 
 MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                         MlirStringRef name) {
   return wrap(unwrap(context)->getOrLoadDialect(unwrap(name)));
 }
 
+bool mlirContextIsRegisteredOperation(MlirContext context, MlirStringRef name) {
+  return unwrap(context)->isOperationRegistered(unwrap(name));
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirDialectGetContext(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getContext());
 }
 

mlir/test/Bindings/Python/dialects.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
+
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testDialectDescriptor
+@run
 def testDialectDescriptor():
   ctx = Context()
   d = ctx.get_dialect_descriptor("std")
   # CHECK: <DialectDescriptor std>
   print(d)
   # CHECK: std
   print(d.namespace)
   try:
     _ = ctx.get_dialect_descriptor("not_existing")
   except ValueError:
     pass
   else:
     assert False, "Expected exception"
 
-run(testDialectDescriptor)
-
 
 # CHECK-LABEL: TEST: testUserDialectClass
+@run
 def testUserDialectClass():
   ctx = Context()
   # Access using attribute.
   d = ctx.dialects.std
   # Note that the standard dialect namespace prints as ''. Others will print
   # as "<Dialect %namespace (..."
   # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
def testUserDialectClass():
   else:
     assert False, "Expected exception"
 
   # Using the 'd' alias.
   d = ctx.d["std"]
   # CHECK: <Dialect (class mlir.dialects._std_ops_gen._Dialect)>
   print(d)
 
-run(testUserDialectClass)
-
 
 # CHECK-LABEL: TEST: testCustomOpView
 # This test uses the standard dialect AddFOp as an example of a user op.
 # TODO: Op creation and access is still quite verbose: simplify this test as
 # additional capabilities come online.
+@run
 def testCustomOpView():
+
   def createInput():
     op = Operation.create("pytest_dummy.intinput", results=[f32])
     # TODO: Auto result cast from operation
     return op.results[0]
 
   with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
def testCustomOpView():
 
   # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[R0:.*]] = addf %[[INPUT0]], %[[INPUT1]] : f32
   # CHECK: %[[R1:.*]] = addf %[[INPUT0]], %[[R0]] : f32
   m.operation.print()
 
 
-run(testCustomOpView)
+# CHECK-LABEL: TEST: testIsRegisteredOperation
+@run
+def testIsRegisteredOperation():
+  ctx = Context()
+
+  # CHECK: std.cond_br: True
+  print(f"std.cond_br: {ctx.is_registered_operation('std.cond_br')}")
+  # CHECK: std.not_existing: False
+  print(f"std.not_existing: {ctx.is_registered_operation('std.not_existing')}")

mlir/test/CAPI/ir.c

int registerOnlyStd() {
   MlirStringRef alsoStdNs = mlirDialectHandleGetNamespace(stdHandle);
   if (stdNs.length != alsoStdNs.length ||
       strncmp(stdNs.data, alsoStdNs.data, stdNs.length))
     return 5;
 
   fprintf(stderr, "@registration\n");
   // CHECK-LABEL: @registration
 
+  // CHECK: std.cond_br is_registered: 1
+  fprintf(stderr, "std.cond_br is_registered: %d\n",
+          mlirContextIsRegisteredOperation(
+              ctx, mlirStringRefCreateFromCString("std.cond_br")));
+
+  // CHECK: std.not_existing_op is_registered: 0
+  fprintf(stderr, "std.not_existing_op is_registered: %d\n",
+          mlirContextIsRegisteredOperation(
+              ctx, mlirStringRefCreateFromCString("std.not_existing_op")));
+
+  // CHECK: not_existing_dialect.not_existing_op is_registered: 0
+  fprintf(stderr, "not_existing_dialect.not_existing_op is_registered: %d\n",
+          mlirContextIsRegisteredOperation(
+              ctx, mlirStringRefCreateFromCString(
+                       "not_existing_dialect.not_existing_op")));
+
   return 0;
 }
 
 /// Tests backreference APIs
 static int testBackreferences() {
   fprintf(stderr, "@test_backreferences\n");
 
   MlirContext ctx = mlirContextCreate();