move append_dialect_search_prefix test

Author: makslevental

mlir/lib/Bindings/Python/MainModule.cpp

@@ -36,6 +36,12 @@ PYBIND11_MODULE(_mlir, m) {
             self.getDialectSearchPrefixes().push_back(std::move(moduleName));
           },
           "module_name"_a)
+      .def(
+          "_check_dialect_module_loaded",
+          [](PyGlobals &self, const std::string &dialectNamespace) {
+            return self.loadDialectModule(dialectNamespace);
+          },
+          "dialect_namespace"_a)
       .def("_register_dialect_impl", &PyGlobals::registerDialectImpl,
            "dialect_namespace"_a, "dialect_class"_a,
            "Testing hook for directly registering a dialect")

mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi

@@ -7,6 +7,7 @@ class _Globals:
     def _register_dialect_impl(self, dialect_namespace: str, dialect_class: type) -> None: ...
     def _register_operation_impl(self, operation_name: str, operation_class: type) -> None: ...
     def append_dialect_search_prefix(self, module_name: str) -> None: ...
+    def _check_dialect_module_loaded(self, dialect_namespace: str) -> bool: ...
 
 def register_dialect(dialect_class: type) -> object: ...
 def register_operation(dialect_class: type) -> object: ...

mlir/test/python/ir/custom_dialect/custom.py

@@ -0,0 +1,4 @@
+# The purpose of this empty dialect module is to enable successfully loading the "custom" dialect.
+# Without this file here (and a corresponding _cext.globals.append_dialect_search_prefix("custom_dialect")),
+# PyGlobals::loadDialectModule would search and fail to find the "custom" dialect for each Operation.create("custom.op")
+# (amongst other things).

mlir/test/python/ir/dialects.py

@@ -1,7 +1,9 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
+import sys
 from mlir.ir import *
+from mlir.dialects._ods_common import _cext
 
 
 def run(f):
@@ -104,3 +106,18 @@ def testIsRegisteredOperation():
     print(f"cf.cond_br: {ctx.is_registered_operation('cf.cond_br')}")
     # CHECK: func.not_existing: False
     print(f"func.not_existing: {ctx.is_registered_operation('func.not_existing')}")
+
+
+# CHECK-LABEL: TEST: testAppendPrefixSearchPath
+@run
+def testAppendPrefixSearchPath():
+    ctx = Context()
+    ctx.allow_unregistered_dialects = True
+    with Location.unknown(ctx):
+        assert not _cext.globals._check_dialect_module_loaded("custom")
+        Operation.create("custom.op")
+        assert not _cext.globals._check_dialect_module_loaded("custom")
+
+        sys.path.append(".")
+        _cext.globals.append_dialect_search_prefix("custom_dialect")
+        assert _cext.globals._check_dialect_module_loaded("custom")

mlir/test/python/ir/insertion_point.py

@@ -1,15 +1,7 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
-import io
-import itertools
-import sys
-
 from mlir.ir import *
-from mlir.dialects._ods_common import _cext
-
-sys.path.append(".")
-_cext.globals.append_dialect_search_prefix("custom_dialect")
 
 
 def run(f):