Migrate experimental MLIR backend (#346)

## Summary

This is a curated migration of the experimental `numbast-mlir` backend
into public Numbast under `numbast.experimental.mlir`, based on public
`main`.

- Adds the experimental MLIR backend package and its source-tree tests
under `numbast/src/numbast/experimental/mlir`.
- Keeps public top-level `numbast` APIs unchanged; MLIR APIs are
imported from `numbast.experimental.mlir` and are not re-exported from
top-level `numbast`.
- Keeps `python -m numbast` as the static binding CLI entrypoint and
routes `MLIR Backend: true` configs to the experimental MLIR static
generator.
- Extends the static binding config schema/docs with `MLIR Backend` and
`Module Link Variables Used`.
- Rejects `Module Link Variables Used` unless `MLIR Backend: true`.
- Excludes the MLIR test subtree from existing broad CI/test commands.

## CI and test coverage

Existing CI environments intentionally do not install or run
`numba_cuda_mlir`. Dedicated `numbast-mlir` CI and real MLIR test
coverage should be added in a follow-up PR.

Current-CI-safe coverage added here:

- Config parsing/routing test that monkeypatches MLIR dispatch without
importing or running `numba_cuda_mlir`.
- Guard test for `Module Link Variables Used` requiring `MLIR Backend:
true`.
- Schema-doc tests for the new config keys.

## Validation

- `PYTHONPATH=numbast/src:ast_canopy python -m pytest
numbast/src/numbast/tools/tests/test_mlir_backend_routing.py
numbast/src/numbast/tools/tests/test_config_schema_docs.py`
- `python -m py_compile ci/run_tests.py
numbast/src/numbast/tools/static_binding_generator.py
numbast/src/numbast/tools/tests/test_mlir_backend_routing.py
numbast/src/numbast/tools/tests/test_config_schema_docs.py`
- `PYTHONPATH=numbast/src:ast_canopy python -m pytest --collect-only
numbast/ --ignore=numbast/src/numbast/experimental/mlir
--ignore=numbast/src/numbast/tools/tests/test_symbol_exposure.py`

The collect-only command needs the extra `test_symbol_exposure.py`
ignore in this local environment because no CUDA device is available;
the MLIR ignore itself prevented `numba_cuda_mlir` collection errors.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Adds an experimental MLIR backend for CUDA/C++ bindings (functions,
function templates, class templates, enums, structs) with configurable
argument-intent semantics and shim/link support.
* Static binding generator now supports an MLIR backend mode and
CLI-driven generation for MLIR workflows.

* **Documentation**
  * Docs updated with MLIR backend requirements and usage notes.

* **Tests / Chores**
* Large suite of MLIR-focused tests added; CI/test configs updated to
skip the experimental MLIR subtree by default.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/numbast/pull/346?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Michael Wang <isVoid@users.noreply.github.com>

Author: isVoid

.github/workflows/wheels-test.yaml

@@ -136,7 +136,9 @@ jobs:
           python -m pytest ast_canopy/tests
       - name: Test numbast wheels
         run: |
-          python -m pytest numbast/
+          python -m pytest \
+            numbast/ \
+            --ignore=numbast/src/numbast/experimental/mlir
       - name: Run numbast_extensions third-party CCCL tests
         run: |
           python -m pytest numbast_extensions/tests/thirdparty/CCCL/

ci/run_tests.py

@@ -5,19 +5,24 @@
 import pytest
 
 
-def run_pytest(lib, test_dir):
+MLIR_TESTS_DIR = "numbast/src/numbast/experimental/mlir"
+
+
+def run_pytest(lib, test_dir, extra_pytest_args=None):
     if RAPIDS_TESTS_DIR := os.environ.get("RAPIDS_TESTS_DIR", None):
         junitxml = os.path.join(RAPIDS_TESTS_DIR, f"junit-{lib}.xml")
     else:
         junitxml = "/dev/null"
 
+    extra_pytest_args = extra_pytest_args or []
     command = [
         "pytest",
         "-v",
         "-s",
         "--continue-on-collection-errors",
         "--cache-clear",
         f"--junitxml={junitxml}",
+        *extra_pytest_args,
         *test_dir,
     ]
     try:
@@ -57,7 +62,7 @@ def run(
     if all_tests or ast_canopy:
         run_pytest("ast_canopy", ["ast_canopy/"])
     if all_tests or numbast:
-        run_pytest("numbast", ["numbast/"])
+        run_pytest("numbast", ["numbast/"], [f"--ignore={MLIR_TESTS_DIR}"])
     if all_tests or bf16:
         run_pytest(
             "bf16",

docs/source/static.rst

@@ -28,6 +28,11 @@ drift between implementation and documentation.
 For details on ``Function Argument Intents`` semantics and generated signatures,
 see :doc:`/argument_intents`.
 
+The experimental ``numbast-mlir`` backend is selected with
+``MLIR Backend: true`` and is imported from ``numbast.experimental.mlir``. It
+requires a user-provided ``numba_cuda_mlir`` installation; the default Numbast
+test environments do not install or run it yet.
+
 Config example:
 
 .. code-block:: yaml

numbast/src/numbast/experimental/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Experimental Numbast namespaces."""
+
+__all__ = []

numbast/src/numbast/experimental/mlir/README.md

@@ -0,0 +1,8 @@
+# Numbast MLIR Backend
+
+This package contains the experimental `numbast-mlir` backend under the
+`numbast.experimental.mlir` namespace.
+
+It is not imported or re-exported by top-level `numbast`. Runtime use requires a
+user-provided `numba_cuda_mlir` installation, and the repository's default CI
+jobs intentionally do not collect or run these tests yet.

numbast/src/numbast/experimental/mlir/__init__.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from numbast.experimental.mlir.struct import bind_cxx_struct, bind_cxx_structs
+from numbast.experimental.mlir.class_template import (
+    bind_cxx_class_template_specialization,
+    bind_cxx_class_template,
+    bind_cxx_class_templates,
+    clear_concrete_type_caches,
+)
+from numbast.experimental.mlir.function import (
+    bind_cxx_function,
+    bind_cxx_functions,
+)
+from numbast.experimental.mlir.function_template import (
+    bind_cxx_function_template,
+    bind_cxx_function_templates,
+)
+from numbast.experimental.mlir.enum import bind_cxx_enum, bind_cxx_enums
+from numbast.experimental.mlir.shim_writer import (
+    MemoryShimWriter,
+    FileShimWriter,
+)
+
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.version("numbast")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0+unknown"
+
+__all__ = [
+    "__version__",
+    "bind_cxx_enum",
+    "bind_cxx_enums",
+    "bind_cxx_function",
+    "bind_cxx_functions",
+    "bind_cxx_function_template",
+    "bind_cxx_function_templates",
+    "bind_cxx_struct",
+    "bind_cxx_structs",
+    "bind_cxx_class_template_specialization",
+    "bind_cxx_class_template",
+    "bind_cxx_class_templates",
+    "clear_concrete_type_caches",
+    "MemoryShimWriter",
+    "FileShimWriter",
+]

numbast/src/numbast/experimental/mlir/args.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from numba_cuda_mlir.mlir_lowering import MLIRLower
+
+from numba_cuda_mlir._mlir import ir
+from numba_cuda_mlir._mlir.dialects import llvm
+
+
+def prepare_ir_types(
+    builder: MLIRLower,
+    argtys: list[ir.Type],
+    *,
+    pass_ptr_mask: list[bool] | None = None,
+) -> list[ir.Type]:
+    """
+    Prepare LLVM IR types for passing function arguments by reference.
+
+    Given a list of argument IR types, return a parallel list of IR types suitable for an ABI that passes arguments by pointer. For each argument, the builder's get_mlir_type() is used to obtain the value type; if the corresponding entry in pass_ptr_mask is True and that value type is already an llvm.PointerType, that pointer type is preserved, otherwise the value type is wrapped in an opaque LLVM pointer type.
+
+    Parameters:
+        builder (MLIRLower): Lowering helper used to obtain the value type via get_mlir_type().
+        argtys (list[ir.Type]): Argument IR types to prepare.
+        pass_ptr_mask (list[bool] | None): Optional mask the same length as argtys indicating per-argument behavior.
+            If None, all entries are treated as False. When True for an argument and the value type is an llvm.PointerType,
+            the pointer type is passed through unchanged.
+
+    Returns:
+        list[ir.Type]: Prepared IR types where each entry is either a pointer-to-value or an existing pointer type preserved per pass_ptr_mask.
+
+    Raises:
+        ValueError: If pass_ptr_mask is provided and its length does not match len(argtys).
+    """
+    if pass_ptr_mask is None:
+        pass_ptr_mask = [False] * len(argtys)
+
+    if len(pass_ptr_mask) != len(argtys):
+        raise ValueError(
+            f"pass_ptr_mask length ({len(pass_ptr_mask)}) must match argtys length ({len(argtys)})"
+        )
+
+    ir_types: list[ir.Type] = []
+    for argty, passthrough in zip(argtys, pass_ptr_mask):
+        vty = builder.get_mlir_type(argty)
+        if passthrough and isinstance(vty, llvm.PointerType):
+            # Pass pointer-typed values directly (e.g. C++ T& mapped to CPointer(T))
+            ir_types.append(vty)
+        else:
+            # Default ABI: pass pointer-to-value
+            ir_types.append(ir.Type.parse("!llvm.ptr", context=vty.context))
+
+    return ir_types

numbast/src/numbast/experimental/mlir/benchmarks/test_arithmetic.py

@@ -0,0 +1,28 @@
+from numba_cuda_mlir import cuda
+from numba_cuda_mlir.types import float32
+import numpy as np
+
+import pytest
+
+
+@pytest.mark.skip(reason="benchmark is not run by default")
+def test_arithmetic(benchmark):
+    from numbast_extensions.bf16 import nv_bfloat16, get_shims  # noqa: E402
+
+    def bench():
+        @cuda.jit(link=get_shims())
+        def kernel(arith):
+            # Binary Arithmetic Operators
+            a = nv_bfloat16(1.0)
+            b = nv_bfloat16(2.0)
+
+            arith[0] = float32(a + b)
+            arith[1] = float32(a - b)
+            arith[2] = float32(a * b)
+            arith[3] = float32(a / b)
+
+        arith = np.zeros(4, dtype=np.float32)
+
+        kernel[1, 1](arith)
+
+    benchmark(bench)

numbast/src/numbast/experimental/mlir/benchmarks/test_rtc.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from tempfile import NamedTemporaryFile
+
+import pytest
+
+from numba_cuda_mlir import cuda
+
+from ast_canopy import parse_declarations_from_source
+from numbast.experimental.mlir import bind_cxx_functions, MemoryShimWriter
+
+function_template = """
+__device__ int {name}(int a, int b) {{
+    return a + b;
+}}
+"""
+
+
+@pytest.fixture(params=[1, 10, 100, 1000])
+def simulate_header(request):
+    """Create a header file that contains N headers to add two integers together.
+    Used to benchmark the impact of number of headers to kernel launch.
+    """
+    N = request.param
+
+    tmp = NamedTemporaryFile(mode="w", suffix=".cuh", delete=False)
+    functions = [function_template.format(name=f"add{i}") for i in range(N)]
+    tmp.write("\n".join(functions))
+    tmp.flush()
+
+    major, minor = cuda.get_current_device().compute_capability
+    decls = parse_declarations_from_source(
+        tmp.name, [tmp.name], f"sm_{major}{minor}"
+    )
+    shim_writer = MemoryShimWriter(f'#include "{tmp.name}"')
+    adds = bind_cxx_functions(shim_writer, decls.functions)
+
+    yield adds, shim_writer
+
+
+@pytest.mark.skip(reason="benchmark is not run by default")
+def test_rtc(benchmark, simulate_header):
+    def bench():
+        adds, shim_writer = simulate_header
+        add = adds[0]
+
+        @cuda.jit(link=shim_writer.links())
+        def kernel():
+            _ = add(1, 2)
+
+        kernel[1, 1]()
+
+    benchmark(bench)