[python] Add MLIR dialect registration hook

docs/sphinx/examples/plugins/mlir_extension/README.md

@@ -12,35 +12,50 @@ The file layout follows standard MLIR conventions:
 include/Trivial/  - public headers for the dialect + pass
 lib/              - the pure MLIR library
 python/           - the nanobind extension built with the CUDA-Q helpers
+cudaq_mlir_extension/  - pip package stub and entry-point declaration
 ```
 
 The dialect and the `trivial-pass` are defined via TableGen, using the
 `mlir-tblgen` binary shipped in the `cudaq-devel` wheel.
 
-## Building
+## Downstream dialect registration
+
+CUDA-Q discovers out-of-tree MLIR dialects through the `cudaq.mlir_dialects`
+entry point group. This example declares one entry point that forwards to
+`register_dialects` on the package:
+
+```toml
+[project.entry-points."cudaq.mlir_dialects"]
+trivial = "cudaq_mlir_extension:register_dialects"
+```
+
+After installation, importing `cudaq` seeds every `cudaq.mlir.ir.Context` with
+the `trivial` dialect automatically—no explicit registration call is required.
 
-Install the `cudaq-devel` wheel into the active Python environment. If you
-are using non-released versions of CUDA-Q (nightly or custom builds), make sure
-your Python package manager knows where to find the `cudaq-devel` wheel and a compatible
-core `cudaq` wheel.
+## Building
 
-Then configure. CMake will look for a valid installation of CUDA-Q within the
-installed Python packages. Make sure to run this command in the same virtual Python
-environment used to install the `cudaq-devel` wheel, or pass the path to the correct
-Python interpreter explicitly using the `-DPython3_EXECUTABLE` flag.
+Install `cudaq-devel` and the matching `cudaq` runtime wheel into your virtual
+environment, then install this example as a normal pip package:
 
 ```bash
-pip install cudaq-devel
-pip install 'nanobind>=2.12.0,<3'
+pip install cudaq-devel cudaq
+pip install /path/to/examples/plugins/mlir_extension
 
-site=$(python -c 'import site; print(site.getsitepackages()[0])')
+python -c "import cudaq; from cudaq.mlir.ir import Context; \
+  ctx = Context(); assert ctx.dialects['trivial'] is not None"
+```
+
+The pip package uses `scikit-build-core` to drive the existing CMake project and
+registers the `cudaq.mlir_dialects` entry point declared in `pyproject.toml`.
+For development, you can also use `cmake` directly:
+
+```bash
+source /path/to/venv/bin/activate  # detects installed cudaq-devel in the venv
 cmake -S examples/plugins/mlir_extension -B build \
   -DCMAKE_BUILD_TYPE=Release \
   -G Ninja
 cmake --build build
-cmake --install build --prefix "$site" --component TrivialMLIRPythonModules
-
-python -c "import sys, glob; \
-  d=glob.glob('$site/cudaq_mlir_extension/mlir/_mlir_libs')[0]; sys.path.insert(0, d); \
-  import _mlirExtension; assert _mlirExtension.run_trivial_pass()"
 ```
+
+although `pip` or equivalent Python tooling will be required to register the
+`cudaq.mlir_dialects` entry point.

docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/__init__.py

@@ -0,0 +1,15 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+"""Minimal out-of-tree CUDA-Q MLIR extension example package."""
+
+
+def register_dialects(registry):
+    """Register the example ``trivial`` dialect with a CUDA-Q MLIR registry."""
+    from cudaq_mlir_extension.mlir._mlir_libs import _mlirExtension
+
+    _mlirExtension.register_dialects(registry)

docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/__init__.py

@@ -0,0 +1,7 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #

docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/_mlir_libs/__init__.py

@@ -0,0 +1,7 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #

docs/sphinx/examples/plugins/mlir_extension/pyproject.toml

@@ -0,0 +1,26 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+[build-system]
+requires = ["nanobind>=2.12.0,<3", "scikit-build-core>=0.10"]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "cudaq-mlir-extension-example"
+version = "0.1.0"
+description = "Minimal out-of-tree CUDA-Q MLIR extension used to validate cudaq-devel"
+requires-python = ">=3.10"
+license = "Apache-2.0"
+dependencies = ["cudaq"]
+
+[project.entry-points."cudaq.mlir_dialects"]
+trivial = "cudaq_mlir_extension:register_dialects"
+
+[tool.scikit-build]
+cmake.source-dir = "."
+wheel.packages = ["cudaq_mlir_extension"]
+install.components = ["TrivialMLIRPythonModules"]

docs/sphinx/examples/plugins/mlir_extension/python/CMakeLists.txt

@@ -23,6 +23,14 @@ add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=cudaq_mlir_extension.mlir.")
 
 declare_mlir_python_sources(TrivialMLIRPythonSources)
 
+declare_mlir_python_sources(TrivialMLIRPythonPackage
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../cudaq_mlir_extension"
+  ADD_TO_PARENT TrivialMLIRPythonSources
+  SOURCES
+    __init__.py
+    mlir/__init__.py
+    mlir/_mlir_libs/__init__.py)
+
 declare_mlir_python_extension(TrivialMLIRPythonSources.Extension
   MODULE_NAME _mlirExtension
   ADD_TO_PARENT TrivialMLIRPythonSources

docs/sphinx/examples/plugins/mlir_extension/python/TrivialExtension.cpp

@@ -15,6 +15,8 @@
 #include "Trivial/TrivialDialect.h"
 #include "Trivial/TrivialPasses.h"
 
+#include "mlir/Bindings/Python/NanobindAdaptors.h"
+#include "mlir/CAPI/IR.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
@@ -27,6 +29,10 @@ NB_MODULE(_mlirExtension, m) {
   m.doc() = "Minimal out-of-tree CUDA-Q MLIR Python extension used to validate "
             "the cudaq-devel wheel.";
 
+  m.def("register_dialects", [](MlirDialectRegistry registry) {
+    unwrap(registry)->insert<trivial::TrivialDialect>();
+  });
+
   m.def("run_trivial_pass", []() {
     MLIRContext context;
     context.getOrLoadDialect<trivial::TrivialDialect>();

python/cudaq/mlir/_mlir_libs/_site_initialize_1.py

@@ -0,0 +1,20 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+"""Dialect registration for downstream CUDA-Q extensions.
+
+Packages that define out-of-tree MLIR dialects advertise a callable in the
+``cudaq.mlir_dialects`` entry point group. Each callable is passed the
+DialectRegistry that every CUDA-Q MLIR Context is seeded from.
+"""
+
+from importlib.metadata import entry_points
+
+
+def register_dialects(registry):
+    for ep in entry_points(group="cudaq.mlir_dialects"):
+        ep.load()(registry)

python/extension/CMakeLists.txt

@@ -65,6 +65,11 @@ declare_mlir_python_extension(CUDAQuantumPythonSources.SiteInitialize
     cudaqMLIR
 )
 
+declare_mlir_python_sources(CUDAQuantumPythonSources.SiteInitializeExtensions
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../cudaq/mlir"
+  ADD_TO_PARENT CUDAQuantumPythonSources
+  SOURCES _mlir_libs/_site_initialize_1.py)
+
 declare_mlir_python_extension(CUDAQuantumPythonSources.Extension
   MODULE_NAME _quakeDialects
   ADD_TO_PARENT CUDAQuantumPythonSources

python/extension/PythonMLIRHooks.cpp

@@ -10,6 +10,11 @@
 #include "cudaq_internal/compiler/TracePassInstrumentation.h"
 #include "runtime/cudaq/platform/PythonSignalCheck.h"
 
+#include "mlir-c/Bindings/Python/Interop.h"
+#include "mlir/CAPI/IR.h"
+
+#include <nanobind/nanobind.h>
+
 // FIXME: Declare this in a header file!
 // Forward-declare the Python-aware helper so this translation unit does not
 // pull in headers from python/. The symbol is defined in
@@ -20,6 +25,8 @@ mlir::LogicalResult runPassManagerReleasingGIL(mlir::PassManager &pm,
                                                mlir::Operation *op);
 }
 
+namespace nb = nanobind;
+
 static mlir::LogicalResult pythonRunPassManager(mlir::PassManager &pm,
                                                 mlir::Operation *op) {
   pm.addInstrumentation(std::make_unique<cudaq::TracePassInstrumentation>());
@@ -28,6 +35,21 @@ static mlir::LogicalResult pythonRunPassManager(mlir::PassManager &pm,
   return cudaq::runPassManagerReleasingGIL(pm, op);
 }
 
+static void pythonRegisterDialects(mlir::DialectRegistry &registry) {
+  if (!Py_IsInitialized())
+    return;
+  nb::gil_scoped_acquire gil;
+  nb::object libs = nb::module_::import_("cudaq.mlir._mlir_libs");
+  nb::object pyRegistry = libs.attr("get_dialect_registry")();
+  MlirDialectRegistry handle = mlirPythonCapsuleToDialectRegistry(
+      pyRegistry.attr(MLIR_PYTHON_CAPI_PTR_ATTR).ptr());
+  if (!mlirDialectRegistryIsNull(handle))
+    unwrap(handle)->appendTo(registry);
+}
+
 namespace cudaq_internal::compiler {
-void installPythonMLIRHooks() { setRunPassManagerHook(&pythonRunPassManager); }
+void installPythonMLIRHooks() {
+  setRunPassManagerHook(&pythonRunPassManager);
+  setDialectRegistrationHook(&pythonRegisterDialects);
+}
 } // namespace cudaq_internal::compiler

python/tests/mlir/test_downstream_dialect_registration.py

@@ -0,0 +1,32 @@
+# ============================================================================ #
+# Copyright (c) 2026 NVIDIA Corporation & Affiliates.                          #
+# All rights reserved.                                                         #
+#                                                                              #
+# This source code and the accompanying materials are made available under     #
+# the terms of the Apache License 2.0 which accompanies this distribution.     #
+# ============================================================================ #
+"""Tests for downstream MLIR dialect registration via entry points."""
+
+# RUN: PYTHONPATH=../../ pytest -rP %s
+
+from unittest.mock import MagicMock, patch
+
+from cudaq.mlir._mlir_libs import _site_initialize_1
+from cudaq.mlir.ir import DialectRegistry
+
+
+def test_site_initialize_dispatches_mlir_dialect_entry_points():
+    registry = DialectRegistry()
+    seen = []
+
+    fake_ep = MagicMock()
+    fake_ep.load.return_value = lambda reg: seen.append(reg)
+
+    with patch.object(_site_initialize_1,
+                      "entry_points",
+                      return_value=[fake_ep]) as mock_eps:
+        _site_initialize_1.register_dialects(registry)
+
+    mock_eps.assert_called_once_with(group="cudaq.mlir_dialects")
+    fake_ep.load.assert_called_once_with()
+    assert seen == [registry]

runtime/internal/compiler/RuntimeMLIR.cpp

@@ -710,6 +710,9 @@ static void registerToIQMJsonTranslation() {
 
 static std::once_flag mlir_init_flag;
 static MLIRContext *mlirContext;
+static cudaq_internal::compiler::DialectRegistrationHook
+    g_dialectRegistrationHook = nullptr;
+
 static std::unique_ptr<MLIRContext> createMLIRContext() {
   // Per-context initialization
   DialectRegistry registry;
@@ -719,6 +722,8 @@ static std::unique_ptr<MLIRContext> createMLIRContext() {
   mlir::LLVM::registerInlinerInterface(registry);
   registerBuiltinDialectTranslation(registry);
   registerLLVMDialectTranslation(registry);
+  if (g_dialectRegistrationHook)
+    g_dialectRegistrationHook(registry);
   auto context = std::make_unique<MLIRContext>(registry);
   context->loadAllAvailableDialects();
   return context;
@@ -808,6 +813,17 @@ void cudaq_internal::compiler::setRunPassManagerHook(RunPassManagerHook hook) {
   g_runPassManagerHook = hook ? hook : &defaultRunPassManager;
 }
 
+void cudaq_internal::compiler::setDialectRegistrationHook(
+    DialectRegistrationHook hook) {
+  g_dialectRegistrationHook = hook;
+  if (!hook || !mlirContext)
+    return;
+  DialectRegistry registry;
+  hook(registry);
+  mlirContext->appendDialectRegistry(registry);
+  mlirContext->loadAllAvailableDialects();
+}
+
 void cudaq_internal::compiler::initializeLangMLIR() {
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();

runtime/internal/compiler/include/cudaq_internal/compiler/RuntimeMLIR.h

@@ -19,6 +19,7 @@
 
 namespace mlir {
 class Dialect;
+class DialectRegistry;
 class MLIRContext;
 class ModuleOp;
 class Operation;
@@ -126,6 +127,11 @@ using RunPassManagerHook = mlir::LogicalResult (*)(mlir::PassManager &,
                                                    mlir::Operation *);
 void setRunPassManagerHook(RunPassManagerHook hook);
 
+/// Hook to register dialects from downstream extensions into the
+/// `MLIRContext`s created by CUDA-Q.
+using DialectRegistrationHook = void (*)(mlir::DialectRegistry &);
+void setDialectRegistrationHook(DialectRegistrationHook hook);
+
 /// Configure the pass manager according to environment variables
 void configurePassManagerFromEnv(mlir::PassManager &pm);
 

scripts/validate_devel_wheel.sh

@@ -196,5 +196,32 @@ then
   exit 1
 fi
 
+echo ""
+echo "PASS: devel wheel validation succeeded (raw extension load)."
+
+echo ""
+echo "=== pip install example package ==="
+if ! pip install -q "$example_src" 2>&1; then
+  echo ""
+  echo "FAIL: pip install of example package failed." >&2
+  exit 1
+fi
+
+echo ""
+echo "=== Verify downstream dialect auto-registration ==="
+if ! python - << 'PY'
+import cudaq
+from cudaq.mlir.ir import Context
+
+with Context() as ctx:
+    _ = ctx.dialects["trivial"]
+print("  OK: trivial dialect auto-registered via cudaq.mlir_dialects entry point")
+PY
+then
+  echo ""
+  echo "FAIL: downstream dialect was not auto-registered in cudaq.mlir.ir.Context." >&2
+  exit 1
+fi
+
 echo ""
 echo "PASS: devel wheel validation succeeded."