Skip to content

Migrate experimental MLIR backend#346

Merged
isVoid merged 6 commits into
mainfrom
codex/migrate-numbast-mlir
May 20, 2026
Merged

Migrate experimental MLIR backend#346
isVoid merged 6 commits into
mainfrom
codex/migrate-numbast-mlir

Conversation

@isVoid
Copy link
Copy Markdown
Collaborator

@isVoid isVoid commented May 20, 2026

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.

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 Change Stack

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 20, 2026

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: ea732e63-6a81-420b-af26-2e87be9a28a1

📥 Commits

Reviewing files that changed from the base of the PR and between 198c9cf and 90b06ff.

📒 Files selected for processing (2)
  • numbast/src/numbast/experimental/mlir/__init__.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_module_link_variables_used.py
💤 Files with no reviewable changes (2)
  • numbast/src/numbast/experimental/mlir/tools/tests/test_module_link_variables_used.py
  • numbast/src/numbast/experimental/mlir/init.py

📝 Walkthrough

Walkthrough

Adds an experimental MLIR backend and static binding generator, MLIR runtime/static modules (types, intents, deduction, callconv, shims), many static/runtime tests and fixtures, package initializers and docs updates, and CI/workflow changes to exclude the experimental MLIR subtree from default test runs.

Changes

Experimental MLIR backend, static generator, and CI/docs wiring

Layer / File(s) Summary
End-to-end MLIR backend, static renderers, CLI, tests, and CI updates
numbast/src/numbast/experimental/..., numbast/src/numbast/tools/..., docs/source/static.rst, .github/workflows/wheels-test.yaml, ci/run_tests.py, pixi.toml
Adds MLIR runtime/static binding modules (types, callconv, deduction, intent, function/class/template binders, shim writers), a static binding generator CLI and MLIR routing, many static/runtime tests and fixtures, package initializers and README, and updates CI/docs to ignore the experimental mlir tests by default.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Poem

A rabbit taps keys by moonlit light,
Shims and enums march into the night,
Tests awake, bindings take flight,
CI skips the grove where experiments bite,
Hops of joy — a small green sprite. 🐇✨

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch codex/migrate-numbast-mlir

@github-actions
Copy link
Copy Markdown

Doc Preview CI
🚀 View pre-built docs at
https://NVIDIA.github.io/numbast/pr-preview/pr-346/

Preview will be ready when GitHub Pages deployment finishes.

@isVoid isVoid marked this pull request as ready for review May 20, 2026 03:21
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 39

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
numbast/src/numbast/experimental/mlir/static/tests/data/operator.cuh (1)

1-16: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick win

Consider adding a header guard to prevent accidental redefinition if this test data file is included multiple times.

This test data file declares struct Foo without a guard (#pragma once or include guards). While test fixtures are typically used in single contexts, adding a guard provides defensive hardening against unintended inclusion.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@numbast/src/numbast/experimental/mlir/static/tests/data/operator.cuh` around
lines 1 - 16, This file defines struct Foo and an overloaded operator+ without
any include guard; add a header guard (e.g., a `#pragma` once or traditional
`#ifndef/define/endif`) around the whole file so that symbols like Foo, Foo::Foo
constructors, and operator+ are not redefined if the file is included multiple
times—place the guard at the top before the struct definition and close it at
the end of the file.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@numbast/src/numbast/experimental/mlir/__init__.py`:
- Around line 6-36: Move the Numba version check (the lines that read
numba.__version__, parse major/minor and raise RuntimeError if int(minor) < 59)
to the top of the module before any imports of MLIR submodules (e.g., before the
imports that bring in bind_cxx_struct, bind_cxx_class_template,
bind_cxx_function, bind_cxx_enum, MemoryShimWriter/FileShimWriter); ensure
__version__ is still set via importlib.metadata as shown and keep the same
RuntimeError message, so the numba version gate runs before importing any
numbast.experimental.mlir.* symbols.

In `@numbast/src/numbast/experimental/mlir/benchmarks/test_arithmetic.py`:
- Around line 12-28: The kernel is being JIT-compiled inside bench() which skews
timing; move the `@cuda.jit`(link=get_shims()) decorated kernel definition out of
the bench() function (keep the same function name and signature, e.g.,
kernel(arith)) so compilation happens once at import-time, then in the benchmark
function call a warmup invocation like kernel[1,1](arith) before the timed call
to benchmark() to prime the JIT and ensure benchmark(bench) measures only
execution.

In `@numbast/src/numbast/experimental/mlir/class_template.py`:
- Around line 848-857: The function bind_cxx_class_template_specialization
currently uses a mutable default aliases={} which can be shared across calls;
change the signature to use aliases: dict[str, list[str]] | None = None (i.e.,
default None) and at the start of bind_cxx_class_template_specialization set
aliases = {} if aliases is None so each call gets a fresh dict; update any
callers or type hints accordingly and ensure you reference the aliases parameter
in that function body when making the replacement.

In `@numbast/src/numbast/experimental/mlir/deduction.py`:
- Around line 161-164: _replace_placeholders currently iterates replacements in
arbitrary order which can corrupt longer placeholders when shorter ones are
replaced first; update _replace_placeholders to iterate the replacement keys
sorted by descending length (use sorted(replacements.keys(), key=len,
reverse=True)) so longer placeholders like "TT" are replaced before shorter ones
like "T", then apply .replace(key, value) for each key to preserve correct
mappings.

In `@numbast/src/numbast/experimental/mlir/function.py`:
- Around line 266-274: The parameter exclude uses a mutable default set() in
bind_cxx_function; change the signature to use None (exclude: set[str] | None =
None) and at the top of bind_cxx_function set exclude = set() if exclude is
None; make the identical change for bind_cxx_functions (update its parameter to
None and initialize it inside the function) so each call gets a fresh set
instead of a shared mutable default.
- Around line 57-124: The code uses assert py_op is not None inside
bind_cxx_operator_overload_function which is skipped under -O; replace it with
an explicit validation that raises a clear exception (e.g., ValueError or
TypeError) when py_op is None. Locate the assert in
bind_cxx_operator_overload_function (the py_op variable derived from
func_decl.overloaded_operator_to_python_operator) and replace it with a
conditional that raises an informative error mentioning the function/mangled
name or operator, before proceeding to register_global(py_op) and lower(py_op,
*param_types). Ensure the new check prevents proceeding when py_op is missing.

In `@numbast/src/numbast/experimental/mlir/functor.py`:
- Around line 57-62: The current name property builds ABI names from
self.func.__name__ + c_params(), which can collide for same-named callables;
change the name property to incorporate a collision-safe qualifier by including
self.func.__module__ and self.func.__qualname__ and appending a short stable
hash (e.g., first 8 chars of sha1) computed from a stable identity string
(module + qualname + function source when available, falling back to
repr(self.func) if not) so ABI names become unique; keep the existing c_params()
normalization and ensure spaces are replaced with underscores as before.
- Around line 44-52: shim() currently assumes two parameters by indexing
self.sig.args[0] and [1], which will crash or produce invalid output for
non-binary functors; update shim (and use functor_shim/N2C/self.name references)
to first validate arity (e.g., check len(self.sig.args) == 2) and either raise a
clear exception or handle other arities (e.g., generate appropriate shim or skip
generation), and ensure RT and T are derived safely from self.sig.return_type
and args only after the arity check so no out-of-range indexing occurs.

In `@numbast/src/numbast/experimental/mlir/overload_selection.py`:
- Around line 50-52: The except block in compute_intent_plan currently catches
all exceptions (except Exception) and appends them to intent_errors; change that
to only catch the documented validation errors raised by compute_intent_plan
(ValueError and TypeError) so unexpected bugs propagate—replace "except
Exception as exc:" with "except (ValueError, TypeError) as exc:" (keeping the
existing intent_errors.append(exc) and continue behavior) in the overload
selection logic where intent_errors is recorded.

In `@numbast/src/numbast/experimental/mlir/shim_writer.py`:
- Around line 69-75: The deduplication currently checks `id` but stores PTX
under the temp filename (`tmp.name`) in `self.ptx_written`, so repeated calls
with the same `id` still write duplicates; update write_to_ptx_shim to key the
map by `id` instead: if `id` not in self.ptx_written then create the
NamedTemporaryFile, write `ptx` to it, and assign self.ptx_written[id] =
tmp.name (or the ptx if you prefer to store contents) so future calls use the
`id` lookup to skip re-writing; ensure references to `self.ptx_written`
elsewhere expect the new key (id) rather than temp filename.
- Around line 58-60: The NamedTemporaryFile handles opened in the constructor
and in write_to_ptx_shim are never closed (leaking descriptors) and the ptx
deduplication is broken because write_to_ptx_shim stores PTX under tmp.name but
checks membership by id; update the code so the temp file objects are closed
after writing (call close() or use a context manager) and change
self.ptx_written to use the provided id as the dict key with the temp file path
(or a small structure { "path": tmp.name }) as the value so membership checks
like `if id not in self.ptx_written` work; also update cleanup logic that
iterates self.ptx_written to remove files by reading the stored paths (e.g.,
from values() or value["path"]) and not from keys once id is the key, and mirror
the same keying approach used by InMemoryShimWriter.

In `@numbast/src/numbast/experimental/mlir/static/enum.py`:
- Around line 17-20: The module currently attaches a FileHandler to file_logger
at import (using getLogger, FileHandler, and logger_path) which causes hardcoded
temp-file writes and handler accumulation on reimport; remove that import-time
FileHandler and instead add file_logger.addHandler(logging.NullHandler()) by
default, and provide a dedicated function (e.g., configure_file_logging(path) or
enable_file_logging(path)) that creates and attaches a FileHandler only when
explicitly called (ensure you check file_logger.handlers to avoid adding
duplicates and avoid hardcoded tempfile.gettempdir() inside the module).

In `@numbast/src/numbast/experimental/mlir/static/function.py`:
- Around line 36-40: The module-level side effects (file_logger, logger_path,
FileHandler installation, and print) must be removed from import-time; instead
add a guarded initializer function (e.g., enable_function_debug() or
configure_function_logger()) that creates logger_path, adds a FileHandler to
file_logger only if a debug flag is set and only when the logger has no existing
file handler, and prints the message when enabled; locate references to
file_logger, logger_path, getLogger, and FileHandler in this file and move the
handler creation/print into that initializer and add a check to avoid re-adding
handlers (inspect file_logger.handlers and only add if none match the
FileHandler path).

In `@numbast/src/numbast/experimental/mlir/static/renderer.py`:
- Around line 49-50: The __init__ of _KeyedStringIO is incorrectly forwarding
kwargs as *kwarg which unpacks dict keys as positional args; update the
signature and forwarding to use standard plural names and proper unpacking so
keyword arguments are passed through (change __init__(self, *arg, **kwarg) to
use *args and **kwargs and call super().__init__(*args, **kwargs)), ensuring
StringIO receives keyword parameters (e.g., initial_value) correctly via
_KeyedStringIO.__init__ and the super().__init__ invocation.

In `@numbast/src/numbast/experimental/mlir/static/struct.py`:
- Around line 1688-1690: The shim include generation in the with_shim_stream
block blindly uses self._default_header (shim_include =
f'"`#include`<{self._default_header}>"') which can be None or incorrect for
per-struct specs; change this to build the include from the actual collected
headers used for the struct (or validate self._default_header before using it).
Locate the with_shim_stream branch and replace the shim_include construction
with logic that: 1) checks if self._default_header is truthy and valid, 2)
otherwise picks the correct header from the struct-specific collected headers
list (the same source used elsewhere when assembling headers), and 3) only call
get_shim(shim_include) when you have a valid include string; keep get_shim and
the rest of the block unchanged.
- Around line 44-48: The module currently creates and installs a FileHandler at
import time via file_logger = getLogger(f"{__name__}"), logger_path =
os.path.join(tempfile.gettempdir(), "test.py") and
file_logger.addHandler(FileHandler(logger_path)); change this so the handler is
only created when explicitly requested (e.g., provide a function like
enable_struct_logging or honor an env var/boolean DEBUG flag) and guard against
duplicate attachments by checking existing handlers on file_logger (e.g.,
inspect file_logger.handlers or use logger.hasHandlers()) before adding a new
FileHandler; ensure the debug message about logger_path is emitted only when the
handler is actually enabled.

In `@numbast/src/numbast/experimental/mlir/static/tests/data/enum.cuh`:
- Around line 6-7: This file uses size_t and memcpy but only includes
<cuda/std/cstdint>, creating a fragile transitive dependency; add direct
includes for <cstddef> (for size_t) and <cstring> (for memcpy) near the top of
the file so symbols used in functions and tests (size_t at lines referenced and
memcpy calls) are guaranteed to be available regardless of other headers.

In
`@numbast/src/numbast/experimental/mlir/static/tests/test_conflict_shim_names.py`:
- Around line 43-47: Replace the hardcoded "/tmp/conflict_func_name_1.py" and
"/tmp/conflict_func_name_2.py" in test_conflict_shim_names.py with files created
under the pytest tmp_path fixture to avoid cross-test collisions; update the two
open(...) write calls to use tmp_path / "conflict_func_name_1.py" and tmp_path /
"conflict_func_name_2.py" (or tmp_path.joinpath(...)) so res1["src"] and
res2["src"] are written to isolated temporary files for each test run.

In
`@numbast/src/numbast/experimental/mlir/static/tests/test_function_static_bindings.py`:
- Around line 112-127: The test_same_argument_types_and_overload parametrization
never uses the dtype value so it never exercises the float overload; update the
kernel to pick the argument constructor based on dtype (e.g., select ctor =
int32 if dtype == "int32" else float32) and call add(ctor(1), ctor(2)) instead
of hardcoding int32, ensuring float32 is imported/available and that add, decl,
and kernel remain the same symbols used in the test.

In `@numbast/src/numbast/experimental/mlir/static/tests/test_link_two_files.py`:
- Around line 77-84: The test kernels foo_kernel and add_kernel create
temporaries (Foo() and add(...)) whose results are unused and may be optimized
away; make their results observable so the optimizer cannot DCE them: in
foo_kernel use the Foo instance in a visible side-effect (e.g., store a field or
return it via a global/results array) and in add_kernel capture the
add(int32(1), int32(2)) result into an observable location (e.g., write into a
device/global buffer or return/append it), ensuring these changes are applied to
the `@cuda.jit`(link=impl) functions foo_kernel and add_kernel so cross-unit
linkage is actually exercised.

In `@numbast/src/numbast/experimental/mlir/static/tests/test_static_demo.py`:
- Around line 26-27: The test currently writes to the fixed absolute path
"/tmp/bindings.py"; change it to use pytest's temporary directory fixture (e.g.,
add a tmp_path parameter to the test function) and write to tmp_path /
"bindings.py" instead (use (tmp_path / "bindings.py").write_text(res["src"]) or
open(tmp_path / "bindings.py", "w")). Also update any subsequent reads/uses of
"/tmp/bindings.py" in this test to reference the tmp_path file object so
parallel runs and portability are safe.

In `@numbast/src/numbast/experimental/mlir/static/types.py`:
- Around line 107-111: The pointer parsing in to_numba_type_str (or the block
handling ty.endswith("*")) uses ty.rstrip("*") which removes all trailing '*'
and loses pointer depth; change the trimming to remove only the last character
(e.g., use ty[:-1].rstrip() or equivalent) so the recursive call to
to_numba_type_str on base_ty correctly produces nested CPointer(...) types, and
keep the BaseRenderer._try_import_numba_type("CPointer") import call as-is.

In `@numbast/src/numbast/experimental/mlir/struct.py`:
- Around line 572-580: The code in bind_cxx_structs (the block using s.name,
aliases, parent_types, data_models) indexes optional maps (parent_types and
data_models) directly which can raise KeyError when callers pass the default {};
change those direct lookups to safe retrievals (e.g., use dict.get with alias or
s.name and handle None) so optional maps are supported — for example, obtain
alias = aliases.get(s.name, [None])[0] when needed, then use
parent_types.get(alias or s.name) and data_models.get(alias or s.name) and
handle the missing-case explicitly (raise a clear ValueError or skip binding)
instead of letting a KeyError propagate. Ensure the logic for unnamed names
still falls back to aliases safely and produces a deterministic error message
when a required spec is absent.

In `@numbast/src/numbast/experimental/mlir/tests/data/sample_enum.cuh`:
- Around line 10-39: The switch statements in device functions eat and
test_cudaRoundMode can leave out[0] uninitialized for unexpected enum values;
update both switch blocks (in eat and test_cudaRoundMode) to handle the default
path by adding a default case that writes a defined sentinel (e.g., 0 or -1)
into out[0] and/or invokes a debug/assert/log to flag the unexpected value so
the device output is always deterministic; ensure the default case is the last
case in each switch so no path leaves out[0] unset.

In `@numbast/src/numbast/experimental/mlir/tests/data/sample_struct.cuh`:
- Line 16: The header defines __device__ void print() which calls printf but
does not include a header declaring printf; add an explicit include for the C
printf declaration (e.g., `#include` <cstdio> or `#include` <stdio.h>) at the top of
the header so the call in print() is declared when the file is compiled in
isolation.

In `@numbast/src/numbast/experimental/mlir/tests/test_class_template.py`:
- Line 23: Update the import of TypingError in test_class_template.py to use the
canonical numba package: replace the current import of TypingError from
numba_cuda_mlir.errors with an import from numba.core.errors so the test uses
numba.core.errors.TypingError for consistency with the rest of the codebase and
existing exception handling.

In `@numbast/src/numbast/experimental/mlir/tools/static_binding_generator.py`:
- Line 46: The code currently registers the custom YAML tag using the global
yaml.Loader which is unsafe; replace that by defining a SafeLoader subclass
(e.g. class NumbastSafeLoader(yaml.SafeLoader): pass) and register the custom
constructor with that loader (use yaml.add_constructor("!numbast_join",
string_constructor, Loader=NumbastSafeLoader) or yaml.SafeLoader.add_constructor
equivalent) and update all yaml.load calls in this module (places referencing
yaml.Loader or plain yaml.load) to use yaml.load(..., Loader=NumbastSafeLoader)
or yaml.load(..., Loader=yaml.SafeLoader) so that string_constructor is only
used with a SafeLoader subclass; target occurrences: the string_constructor
registration lines and any yaml.load invocation in static_binding_generator.py.
- Around line 782-783: The --cfg-path Click option is optional but cfg_path is
used unconditionally by Config.from_yaml_path(cfg_path), causing a runtime
TypeError when None is passed; update the Click option declaration for
"--cfg-path" (the parameter cfg_path) to include required=True (same pattern as
the --output-dir option) so Click enforces the presence of cfg_path before
execution.
- Around line 714-716: The relative path calculation for config_rel_path is
using output_file (a file path) as the start for os.path.relpath which yields an
incorrect ../ count; change the start to the containing directory of the output
file—e.g., compute output_dir = os.path.dirname(output_file) (or
os.path.abspath(os.path.dirname(...)) if needed) and call
os.path.relpath(cfg_file_path, output_dir) so config_rel_path is relative to the
output directory; update the code around cfg_file_path, config_rel_path, and the
os.path.relpath call accordingly.

In
`@numbast/src/numbast/experimental/mlir/tools/tests/test_bypass_parsing_error.py`:
- Around line 24-30: The test calls
run_in_isolated_folder("bypass_parse_errors.yml.j2", "error_code.cuh",
{"arch_str": arch_str}, ruff_format=False,
bypass_parse_error=bypass_parse_error) but never asserts the bypass-enabled path
succeeded; update the test to capture the result from run_in_isolated_folder (or
check for the expected generated artifact) and add an explicit assertion that it
succeeded — e.g., assign result = run_in_isolated_folder(...) and assert result
is truthy or assert the expected output file (like "error_code.cuh")
exists/contains expected content when bypass_parse_error is True; reference the
run_in_isolated_folder call and bypass_parse_error parameter when adding the
assertion.

In `@numbast/src/numbast/experimental/mlir/tools/tests/test_cli.py`:
- Around line 82-83: The test is currently masking failures by wrapping an
assert inside pytest.raises(Exception); replace that pattern by removing the
pytest.raises context and directly assert the command returned a non-zero exit
using the existing result variable (i.e., check result.exit_code != 0) so the
test explicitly verifies a failing exit instead of catching an AssertionError;
remove the pytest.raises(Exception) usage around the assert to fix test_cli.py.
- Line 11: Replace the import of TypingError from numba_cuda_mlir.errors with
the standard repo pattern by importing TypingError from numba.core.errors;
update the import statement that currently references TypingError to use
numba.core.errors to ensure cross-version compatibility and consistency with
other modules that rely on TypingError.

In `@numbast/src/numbast/experimental/mlir/tools/tests/test_module_callbacks.py`:
- Around line 53-58: The subprocess.run calls that execute test kernels (the
call assigning to res using [sys.executable, test_kernel] with cwd=output_folder
and stdout/stderr capture) need a timeout to avoid CI hangs; add a timeout
argument (e.g. timeout=30 or a test-configurable value) to each subprocess.run
(locations of the three occurrences around the current call and at the other two
ranges) and catch subprocess.TimeoutExpired to terminate/cleanup and fail the
test while including the captured output in the failure message so the test
fails deterministically instead of blocking.

In
`@numbast/src/numbast/experimental/mlir/tools/tests/test_module_link_variables_used.py`:
- Line 47: The test is eagerly evaluating linker.variables_used when calling
getattr so an AttributeError may be raised; change the getattr call to use None
as the default (e.g., getattr(linker, "variable_used", None)), then if that
result is None fall back to linker.variables_used (or assign variables_used to a
local and compare) before asserting equality — locate uses of getattr(linker,
"variable_used", linker.variables_used) in the test and replace with the
two-step safe check using None as the default and an explicit fallback to
linker.variables_used.

In `@numbast/src/numbast/experimental/mlir/tools/tests/test_output_name.py`:
- Around line 47-52: The subprocess.run call that executes the test kernel
(variable test_kernel, invoked into cwd=output_folder and assigned to res) needs
a timeout so the test cannot hang; add a timeout parameter (e.g. timeout=30 or a
CI-configurable value) to subprocess.run and handle subprocess.TimeoutExpired
around that call in the test (fail the test and include any partial
output/stderr in the failure message) so the CI job fails fast instead of
hanging.

In
`@numbast/src/numbast/experimental/mlir/tools/tests/test_shim_include_override.py`:
- Around line 32-40: The test currently sets os_is_imported true for any
"import" line and never asserts that a "shim_include" line was actually seen,
allowing false positives; change the loop over bindings to detect specifically
an "import os" (e.g., check line.startswith("import os") or a regex for
importing os) and separately track a found_shim_include flag when a line
startswith("shim_include") and contains "os", then after the loop assert both
os_is_imported and found_shim_include are True (use the existing variables
bindings, os_is_imported and introduce found_shim_include) so the test fails if
the shim include override is not emitted or if os was not imported.

In `@numbast/src/numbast/experimental/mlir/tools/tests/test_symbol_exposure.py`:
- Around line 8-12: The module currently constructs a CUDA Device at import time
via the module-level dev = Device(0) and cc = dev.compute_capability which
blocks pytest collection on machines without CUDA; move the Device(0) probe into
the test function (or a fixture) so the probe runs at test execution time and
wrap it with a try/except that calls pytest.skip on failure (or reuse the
existing arch_str/cc fixture patterns from conftest.py and test_shim_writer.py
as examples). Replace the module-level dev/cc usage with a local probe inside
the test (or accept the cc fixture) and ensure failure paths call pytest.skip
with a clear message.

In `@numbast/src/numbast/tools/static_binding_generator.py`:
- Around line 58-64: The probe function _config_dict_uses_mlir_backend currently
assumes config_dict is a mapping and will call .get on null/scalar/list YAML
roots; guard it by first validating the input: if not isinstance(config_dict,
dict) raise a clear ValueError (e.g. "Invalid config: expected mapping for root,
got <type>") so callers get an explicit config error instead of an
AttributeError, then proceed to read config_dict.get(...). Apply the same
defensive type check and error for the other config-probing helper(s) referenced
around lines 101-104 so all YAML-root checks consistently validate mapping input
before using .get.
- Around line 101-104: The helper _cfg_path_uses_mlir_backend currently calls
yaml.load(cfg_path, yaml.Loader) which is unsafe for untrusted configs; replace
that call with yaml.safe_load(...) to parse the file safely. Also search this
module for other yaml.load usages (the other occurrence in this file) and the
two occurrences in numbast/experimental/mlir/tools/static_binding_generator.py
and update them to yaml.safe_load as well, keeping the same variable names and
behavior otherwise (i.e., read file, parse into config_dict, then call
_config_dict_uses_mlir_backend or the corresponding consumer).

---

Outside diff comments:
In `@numbast/src/numbast/experimental/mlir/static/tests/data/operator.cuh`:
- Around line 1-16: This file defines struct Foo and an overloaded operator+
without any include guard; add a header guard (e.g., a `#pragma` once or
traditional `#ifndef/define/endif`) around the whole file so that symbols like
Foo, Foo::Foo constructors, and operator+ are not redefined if the file is
included multiple times—place the guard at the top before the struct definition
and close it at the end of the file.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: b7625bb6-4f8f-4fe0-8578-d345a0377eb2

📥 Commits

Reviewing files that changed from the base of the PR and between 90d9a90 and 198c9cf.

📒 Files selected for processing (130)
  • .github/workflows/wheels-test.yaml
  • ci/run_tests.py
  • docs/source/static.rst
  • numbast/src/numbast/experimental/__init__.py
  • numbast/src/numbast/experimental/mlir/README.md
  • numbast/src/numbast/experimental/mlir/__init__.py
  • numbast/src/numbast/experimental/mlir/args.py
  • numbast/src/numbast/experimental/mlir/benchmarks/test_arithmetic.py
  • numbast/src/numbast/experimental/mlir/benchmarks/test_rtc.py
  • numbast/src/numbast/experimental/mlir/callconv.py
  • numbast/src/numbast/experimental/mlir/class_template.py
  • numbast/src/numbast/experimental/mlir/deduction.py
  • numbast/src/numbast/experimental/mlir/enum.py
  • numbast/src/numbast/experimental/mlir/errors.py
  • numbast/src/numbast/experimental/mlir/function.py
  • numbast/src/numbast/experimental/mlir/function_template.py
  • numbast/src/numbast/experimental/mlir/functor.py
  • numbast/src/numbast/experimental/mlir/intent.py
  • numbast/src/numbast/experimental/mlir/intent_defs.py
  • numbast/src/numbast/experimental/mlir/overload_selection.py
  • numbast/src/numbast/experimental/mlir/shim_writer.py
  • numbast/src/numbast/experimental/mlir/static/__init__.py
  • numbast/src/numbast/experimental/mlir/static/callconv.py
  • numbast/src/numbast/experimental/mlir/static/enum.py
  • numbast/src/numbast/experimental/mlir/static/function.py
  • numbast/src/numbast/experimental/mlir/static/renderer.py
  • numbast/src/numbast/experimental/mlir/static/struct.py
  • numbast/src/numbast/experimental/mlir/static/tests/conftest.py
  • numbast/src/numbast/experimental/mlir/static/tests/data/bf16.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/conflict_func_name_1.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/conflict_func_name_2.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/demo.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/enum.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/function.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/function_out.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/multiple_decls.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/operator.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/data/src/function.cu
  • numbast/src/numbast/experimental/mlir/static/tests/data/src/function_out.cu
  • numbast/src/numbast/experimental/mlir/static/tests/data/src/operator.cu
  • numbast/src/numbast/experimental/mlir/static/tests/data/src/struct.cu
  • numbast/src/numbast/experimental/mlir/static/tests/data/struct.cuh
  • numbast/src/numbast/experimental/mlir/static/tests/test_bf16_support.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_conflict_shim_names.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_enum_static_bindings.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_function_static_bindings.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_link_two_files.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_multiple_decls.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_operator_bindings.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_static_demo.py
  • numbast/src/numbast/experimental/mlir/static/tests/test_struct_static_bindings.py
  • numbast/src/numbast/experimental/mlir/static/typedef.py
  • numbast/src/numbast/experimental/mlir/static/types.py
  • numbast/src/numbast/experimental/mlir/struct.py
  • numbast/src/numbast/experimental/mlir/tests/conftest.py
  • numbast/src/numbast/experimental/mlir/tests/data/sample_class_template.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_class_template_templated_method.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_enum.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_function.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_function_mutative.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_function_out.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_function_simple.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_function_template.cuh
  • numbast/src/numbast/experimental/mlir/tests/data/sample_struct.cuh
  • numbast/src/numbast/experimental/mlir/tests/demo/demo.cuh
  • numbast/src/numbast/experimental/mlir/tests/demo/demo.py
  • numbast/src/numbast/experimental/mlir/tests/demo/test_demo.py
  • numbast/src/numbast/experimental/mlir/tests/test_class_template.py
  • numbast/src/numbast/experimental/mlir/tests/test_deduction.py
  • numbast/src/numbast/experimental/mlir/tests/test_enum_arg.py
  • numbast/src/numbast/experimental/mlir/tests/test_function.py
  • numbast/src/numbast/experimental/mlir/tests/test_function_simple.py
  • numbast/src/numbast/experimental/mlir/tests/test_function_template.py
  • numbast/src/numbast/experimental/mlir/tests/test_numbast_version.py
  • numbast/src/numbast/experimental/mlir/tests/test_shim_writer.py
  • numbast/src/numbast/experimental/mlir/tests/test_struct.py
  • numbast/src/numbast/experimental/mlir/tools/__init__.py
  • numbast/src/numbast/experimental/mlir/tools/static_binding_generator.py
  • numbast/src/numbast/experimental/mlir/tools/tests/additional_include.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/config/bypass_parse_errors.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/cc.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/cfg.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/cooperative_launch.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/cooperative_launch_invalid_regex.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/cooperative_launch_regex.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/module_callbacks.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/module_link_variables_used.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/output_name.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/predefined_macros.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/prefix_removal.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/shim_include_override.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/skip_prefix.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/config/use_separate_registry.yml.j2
  • numbast/src/numbast/experimental/mlir/tools/tests/conftest.py
  • numbast/src/numbast/experimental/mlir/tools/tests/data.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/data2.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/data_ctor_lowering.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/error_code.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/include/bar.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/macros.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/module_callbacks.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/module_link_variables_used.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/predefined_macros.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/prefix_removal.cuh
  • numbast/src/numbast/experimental/mlir/tools/tests/test_additional_includes.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_bypass_parsing_error.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_cc.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_cli.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_macros.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_module_callbacks.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_module_link_variables_used.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_output_name.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_predefined_macros.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_prefix_removal.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_repro_info.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_shim_include_override.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_skip_prefix.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_symbol_exposure.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_use_cooperative.py
  • numbast/src/numbast/experimental/mlir/tools/tests/test_use_separate_registry.py
  • numbast/src/numbast/experimental/mlir/tools/tests/use_cooperative.cuh
  • numbast/src/numbast/experimental/mlir/tools/yaml_tags.py
  • numbast/src/numbast/experimental/mlir/type_naming.py
  • numbast/src/numbast/experimental/mlir/types.py
  • numbast/src/numbast/experimental/mlir/utils.py
  • numbast/src/numbast/tools/static_binding_generator.py
  • numbast/src/numbast/tools/static_binding_generator.schema.yaml
  • numbast/src/numbast/tools/tests/test_config_schema_docs.py
  • numbast/src/numbast/tools/tests/test_mlir_backend_routing.py
  • pixi.toml

Comment thread numbast/src/numbast/experimental/mlir/__init__.py Outdated
Comment thread numbast/src/numbast/experimental/mlir/benchmarks/test_arithmetic.py
Comment thread numbast/src/numbast/experimental/mlir/class_template.py
Comment thread numbast/src/numbast/experimental/mlir/deduction.py
Comment thread numbast/src/numbast/experimental/mlir/function.py
Comment thread numbast/src/numbast/tools/static_binding_generator.py
Comment thread numbast/src/numbast/tools/static_binding_generator.py
@isVoid
Copy link
Copy Markdown
Collaborator Author

isVoid commented May 20, 2026

Plan: I'm going to merge this as it's been tested locally passing numba-cuda-mlir tests. A follow up to this PR is to create CI jobs which installs numba-cuda-mlir into the environment, we can detect regressions there.

@isVoid isVoid merged commit fb74340 into main May 20, 2026
32 checks passed
@isVoid isVoid mentioned this pull request May 22, 2026
isVoid added a commit that referenced this pull request May 23, 2026
## Summary
- update root `VERSION` to `0.10.0`
- add `0.10.0` to `docs/versions.json` and `docs/nv-versions.json` for
docs version picker support

## Changelog
- Remove direct numba imports from MLIR backend (#365)
- [codex] add numba-cuda-mlir support guide (#364)
- Split CI tests for Numba-CUDA and MLIR (#363)
- Migrate experimental MLIR backend (#346)
- Bump test-summary/action from 2.4 to 2.6 in the actions-monthly group
(#345)
- [codex] Support POD struct array fields (#343)
- fix(types,callconv): register CUDA vector ABI alignment (#321)

## Validation
- `python -m json.tool docs/versions.json`
- `python -m json.tool docs/nv-versions.json`
- `git diff --check`


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

## Summary by CodeRabbit

* **Chores**
  * Version 0.10.0 released
  * Updated version references in documentation configuration files

<!-- 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/367?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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant