diff --git a/docs/sphinx/examples/plugins/mlir_extension/README.md b/docs/sphinx/examples/plugins/mlir_extension/README.md index cc5ed096780..b841f846320 100644 --- a/docs/sphinx/examples/plugins/mlir_extension/README.md +++ b/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. diff --git a/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/__init__.py b/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/__init__.py new file mode 100644 index 00000000000..25f4133af04 --- /dev/null +++ b/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) diff --git a/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/__init__.py b/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/__init__.py new file mode 100644 index 00000000000..c6c6f4d157c --- /dev/null +++ b/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. # +# ============================================================================ # diff --git a/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/_mlir_libs/__init__.py b/docs/sphinx/examples/plugins/mlir_extension/cudaq_mlir_extension/mlir/_mlir_libs/__init__.py new file mode 100644 index 00000000000..c6c6f4d157c --- /dev/null +++ b/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. # +# ============================================================================ # diff --git a/docs/sphinx/examples/plugins/mlir_extension/pyproject.toml b/docs/sphinx/examples/plugins/mlir_extension/pyproject.toml new file mode 100644 index 00000000000..5e6ea30ce43 --- /dev/null +++ b/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"] diff --git a/docs/sphinx/examples/plugins/mlir_extension/python/CMakeLists.txt b/docs/sphinx/examples/plugins/mlir_extension/python/CMakeLists.txt index af462b59f3f..aa9cf9e1b7d 100644 --- a/docs/sphinx/examples/plugins/mlir_extension/python/CMakeLists.txt +++ b/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 diff --git a/docs/sphinx/examples/plugins/mlir_extension/python/TrivialExtension.cpp b/docs/sphinx/examples/plugins/mlir_extension/python/TrivialExtension.cpp index fedc7b359bb..68f9bd4be81 100644 --- a/docs/sphinx/examples/plugins/mlir_extension/python/TrivialExtension.cpp +++ b/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(); + }); + m.def("run_trivial_pass", []() { MLIRContext context; context.getOrLoadDialect(); diff --git a/python/cudaq/mlir/_mlir_libs/_site_initialize_1.py b/python/cudaq/mlir/_mlir_libs/_site_initialize_1.py new file mode 100644 index 00000000000..bb3a12a56ad --- /dev/null +++ b/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) diff --git a/python/extension/CMakeLists.txt b/python/extension/CMakeLists.txt index 1fb6fd4ca59..244399cb5fe 100644 --- a/python/extension/CMakeLists.txt +++ b/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 diff --git a/python/extension/PythonMLIRHooks.cpp b/python/extension/PythonMLIRHooks.cpp index 5ae72b653d5..f1efbfc0ff1 100644 --- a/python/extension/PythonMLIRHooks.cpp +++ b/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 + // 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()); @@ -28,6 +35,21 @@ static mlir::LogicalResult pythonRunPassManager(mlir::PassManager &pm, return cudaq::runPassManagerReleasingGIL(pm, op); } +static void pythonRegisterDialects(mlir::DialectRegistry ®istry) { + 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 diff --git a/python/tests/mlir/test_downstream_dialect_registration.py b/python/tests/mlir/test_downstream_dialect_registration.py new file mode 100644 index 00000000000..1bff7fdc220 --- /dev/null +++ b/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] diff --git a/runtime/internal/compiler/RuntimeMLIR.cpp b/runtime/internal/compiler/RuntimeMLIR.cpp index be926ed4f66..8da12893bba 100644 --- a/runtime/internal/compiler/RuntimeMLIR.cpp +++ b/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 createMLIRContext() { // Per-context initialization DialectRegistry registry; @@ -719,6 +722,8 @@ static std::unique_ptr createMLIRContext() { mlir::LLVM::registerInlinerInterface(registry); registerBuiltinDialectTranslation(registry); registerLLVMDialectTranslation(registry); + if (g_dialectRegistrationHook) + g_dialectRegistrationHook(registry); auto context = std::make_unique(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(); diff --git a/runtime/internal/compiler/include/cudaq_internal/compiler/RuntimeMLIR.h b/runtime/internal/compiler/include/cudaq_internal/compiler/RuntimeMLIR.h index 2cd162e2e59..10516c934a3 100644 --- a/runtime/internal/compiler/include/cudaq_internal/compiler/RuntimeMLIR.h +++ b/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); diff --git a/scripts/validate_devel_wheel.sh b/scripts/validate_devel_wheel.sh index 5cec84a9f29..5641d8218e1 100755 --- a/scripts/validate_devel_wheel.sh +++ b/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."