From 3cac86adfce0bf65ba6d80d36dee0f49b25d8988 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 3 Jun 2026 07:07:25 -0700 Subject: [PATCH 1/2] plumbing --- conda/recipes/cudf/recipe.yaml | 3 + dependencies.yaml | 27 ++++ .../cudf/core/udf/mlir_backend/__init__.py | 11 ++ .../private_objects/mlir_backend/__init__.py | 2 + .../mlir_backend/test_plumbing.py | 119 ++++++++++++++++++ python/cudf/cudf/utils/_numba.py | 9 +- python/cudf/pyproject.toml | 1 + 7 files changed, 171 insertions(+), 1 deletion(-) create mode 100644 python/cudf/cudf/core/udf/mlir_backend/__init__.py create mode 100644 python/cudf/cudf/tests/private_objects/mlir_backend/__init__.py create mode 100644 python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py diff --git a/conda/recipes/cudf/recipe.yaml b/conda/recipes/cudf/recipe.yaml index e0cb0caecba5..ce854981f219 100644 --- a/conda/recipes/cudf/recipe.yaml +++ b/conda/recipes/cudf/recipe.yaml @@ -91,6 +91,9 @@ requirements: - pandas >=3.0.0,<3.1.0 - cupy >=13.6.0,!=14.0.0,!=14.1.0 - numba-cuda >=0.22.2,<0.29.0 + # TODO: add `numba-cuda-mlir >=0.3.0` here once it is published to a conda + # channel. Currently only available via PyPI; the pip-only declaration + # lives in `dependencies.yaml::depends_on_numba_cuda_mlir`. - numba >=0.60.0,<0.65.0 - numpy >=1.26,<3.0 - pyarrow>=19.0.0 diff --git a/dependencies.yaml b/dependencies.yaml index decbaacaef10..54b2cce4ec98 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -24,6 +24,7 @@ files: - depends_on_librmm - depends_on_libnvcomp - depends_on_numba_cuda + - depends_on_numba_cuda_mlir - depends_on_rapids_logger - depends_on_ray - depends_on_rmm @@ -86,6 +87,7 @@ files: - depends_on_pylibcudf - depends_on_libcudf - depends_on_numba_cuda + - depends_on_numba_cuda_mlir test_python_pylibcudf: output: none includes: @@ -198,6 +200,7 @@ files: - depends_on_cupy - depends_on_libcudf - depends_on_numba_cuda + - depends_on_numba_cuda_mlir - depends_on_pylibcudf_pyarrow - depends_on_rmm - depends_on_cuda_python @@ -1095,6 +1098,30 @@ dependencies: - matrix: packages: - *numba_cuda + depends_on_numba_cuda_mlir: + # numba-cuda-mlir is the MLIR-based numba-cuda compiler used by the cudf + # UDF backend. Currently only published to PyPI; conda packaging is a + # tracked feature request (see TODOs in conda/recipes/cudf/recipe.yaml). + # Once a conda package exists, add an `output_types: [conda]` entry here + # mirroring `depends_on_numba_cuda`. + specific: + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + cuda_suffixed: "true" + packages: + - numba-cuda-mlir[cu12]>=0.3.0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - numba-cuda-mlir[cu13]>=0.3.0 + # fallback to numba-cuda-mlir with no extra CUDA packages if + # 'cuda_suffixed' isn't true + - matrix: + packages: + - numba-cuda-mlir>=0.3.0 depends_on_pylibcudf: common: - output_types: conda diff --git a/python/cudf/cudf/core/udf/mlir_backend/__init__.py b/python/cudf/cudf/core/udf/mlir_backend/__init__.py new file mode 100644 index 000000000000..79043c3c2fc0 --- /dev/null +++ b/python/cudf/cudf/core/udf/mlir_backend/__init__.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""MLIR / numba_cuda_mlir UDF backend (scaffolding). + +Future PRs in this series will populate this package with typing/lowering +modules that register `cudf` UDF types (``MaskedType``, ``StringView``, +``GroupType``, etc.) with ``numba_cuda_mlir``. To avoid circular imports +through ``cudf.core.udf.utils``, registration is performed by the parent +``cudf.core.udf`` ``__init__`` once the modules exist; this file is +intentionally empty. +""" diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/__init__.py b/python/cudf/cudf/tests/private_objects/mlir_backend/__init__.py new file mode 100644 index 000000000000..8eca3cc68ad4 --- /dev/null +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py new file mode 100644 index 000000000000..0ed07a52b2bd --- /dev/null +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the numba_cuda_mlir UDF backend scaffolding (PR 1). + +This PR introduces a shell of the future MLIR backend without any code +that depends on ``numba_cuda_mlir``. Tests here cover only the +infrastructure pieces: + +* the ``cudf.core.udf.mlir_backend`` package is importable as an empty + scaffold, +* ``cudf.utils._numba`` defaults ``numba.cuda.config.CUDA_ENABLE_NRT`` to + ``False`` at import time, +* the existing ``nrt_enabled()`` and ``CaptureNRTUsage`` helpers in + ``cudf.core.udf.nrt_utils`` behave as documented. + +Subsequent PRs will add real typing/lowering tests under this directory +(e.g. ``test_strings_typing.py``, ``test_masked_lowering.py``). +""" + +from __future__ import annotations + +import importlib + +import pytest +from numba.cuda import config as numba_config + +from cudf.core.udf.nrt_utils import ( + CaptureNRTUsage, + _current_nrt_context, + nrt_enabled, +) + + +def test_mlir_backend_package_importable(): + """The empty mlir_backend package can be imported without side effects.""" + mod = importlib.import_module("cudf.core.udf.mlir_backend") + assert mod.__doc__ and "MLIR" in mod.__doc__ + assert not hasattr(mod, "MaskedType"), ( + "PR 1 should ship an empty mlir_backend; later PRs add types here." + ) + + +def test_cuda_enable_nrt_default_false(): + """`cudf.utils._numba` should default `CUDA_ENABLE_NRT` to False on import.""" + import cudf.utils._numba # noqa: F401 -- import for its side effect + + assert numba_config.CUDA_ENABLE_NRT is False, ( + "Expected CUDA_ENABLE_NRT False after `cudf.utils._numba` import; " + "kernels that need NRT should opt in via `nrt_enabled()`." + ) + + +def test_nrt_enabled_round_trips_to_false(): + """`nrt_enabled()` flips the global on, then restores it.""" + import cudf.utils._numba # noqa: F401 + + assert numba_config.CUDA_ENABLE_NRT is False + with nrt_enabled(): + assert numba_config.CUDA_ENABLE_NRT is True + assert numba_config.CUDA_ENABLE_NRT is False + + +def test_nrt_enabled_restores_pre_existing_true_value(): + """If a caller manually sets the flag True, exiting the context preserves it.""" + saved = numba_config.CUDA_ENABLE_NRT + try: + numba_config.CUDA_ENABLE_NRT = True + with nrt_enabled(): + assert numba_config.CUDA_ENABLE_NRT is True + # Should restore to True (pre-context value), not to False. + assert numba_config.CUDA_ENABLE_NRT is True + finally: + numba_config.CUDA_ENABLE_NRT = saved + + +def test_nrt_enabled_restores_on_exception(): + """An exception inside `nrt_enabled()` must still restore the prior value.""" + import cudf.utils._numba # noqa: F401 + + assert numba_config.CUDA_ENABLE_NRT is False + with pytest.raises(RuntimeError, match="boom"): + with nrt_enabled(): + assert numba_config.CUDA_ENABLE_NRT is True + raise RuntimeError("boom") + assert numba_config.CUDA_ENABLE_NRT is False + + +def test_capture_nrt_usage_default_false(): + """`CaptureNRTUsage` starts with `use_nrt = False`.""" + cap = CaptureNRTUsage() + assert cap.use_nrt is False + + +def test_capture_nrt_usage_observes_inner_set(): + """Code inside the `with` block can flip `use_nrt` and the captor sees it.""" + with CaptureNRTUsage() as cap: + # A type instantiation that needs NRT would do this internally: + ctx = _current_nrt_context.get(None) + assert ctx is cap + ctx.use_nrt = True + assert cap.use_nrt is True + + +def test_capture_nrt_usage_context_var_unset_outside(): + """The context var is reset on exit so nested code can't observe a stale captor.""" + with CaptureNRTUsage(): + pass + assert _current_nrt_context.get(None) is None + + +def test_capture_nrt_usage_nested(): + """Nesting two captors restores the outer context on inner exit.""" + with CaptureNRTUsage() as outer: + with CaptureNRTUsage() as inner: + assert _current_nrt_context.get(None) is inner + inner.use_nrt = True + assert _current_nrt_context.get(None) is outer + assert outer.use_nrt is False + assert inner.use_nrt is True diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py index 79fdb1b1fc65..8d62f8e7498f 100644 --- a/python/cudf/cudf/utils/_numba.py +++ b/python/cudf/cudf/utils/_numba.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -6,6 +6,13 @@ from numba.cuda import config as numba_config from packaging import version +# Default NRT off for all numba-CUDA compilations driven through cudf. The +# UDF compile path (`cudf.core.udf.nrt_utils.nrt_enabled`) flips this to True +# per-kernel for UDFs whose data model declares an NRT meminfo (e.g. string +# returns), and restores it afterwards. Setting the global default to False +# avoids linking NRT into kernels that don't need it. +numba_config.CUDA_ENABLE_NRT = False + # Avoids using contextlib.contextmanager due to additional overhead class _CUDFNumbaConfig: diff --git a/python/cudf/pyproject.toml b/python/cudf/pyproject.toml index 2f646939f8ed..72dbe2509774 100644 --- a/python/cudf/pyproject.toml +++ b/python/cudf/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", "fsspec>=0.6.0", "libcudf==26.8.*,>=0.0.0a0", + "numba-cuda-mlir>=0.3.0", "numba-cuda>=0.22.2,<0.29.0", "numba>=0.60.0,<0.65.0", "numpy>=1.26,<3.0", From 7f73bb4ab4590bcc0d75d2a00141951727455195 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Wed, 10 Jun 2026 11:52:24 -0700 Subject: [PATCH 2/2] reduce PR scope --- .../mlir_backend/test_plumbing.py | 92 ++----------------- python/cudf/cudf/utils/_numba.py | 9 +- 2 files changed, 7 insertions(+), 94 deletions(-) diff --git a/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py b/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py index a56137aed309..4acd532fe875 100644 --- a/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py +++ b/python/cudf/cudf/tests/private_objects/mlir_backend/test_plumbing.py @@ -1,93 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 - - -from __future__ import annotations - import pytest -from numba.cuda import config as numba_config - -from cudf.core.udf.nrt_utils import ( - CaptureNRTUsage, - _current_nrt_context, - nrt_enabled, -) - - -def test_cuda_enable_nrt_default_false(): - """`cudf.utils._numba` should default `CUDA_ENABLE_NRT` to False on import.""" - import cudf.utils._numba # noqa: F401 -- import for its side effect - - assert numba_config.CUDA_ENABLE_NRT is False, ( - "Expected CUDA_ENABLE_NRT False after `cudf.utils._numba` import; " - "kernels that need NRT should opt in via `nrt_enabled()`." - ) - - -def test_nrt_enabled_round_trips_to_false(): - """`nrt_enabled()` flips the global on, then restores it.""" - import cudf.utils._numba # noqa: F401 - - assert numba_config.CUDA_ENABLE_NRT is False - with nrt_enabled(): - assert numba_config.CUDA_ENABLE_NRT is True - assert numba_config.CUDA_ENABLE_NRT is False - - -def test_nrt_enabled_restores_pre_existing_true_value(): - """If a caller manually sets the flag True, exiting the context preserves it.""" - saved = numba_config.CUDA_ENABLE_NRT - try: - numba_config.CUDA_ENABLE_NRT = True - with nrt_enabled(): - assert numba_config.CUDA_ENABLE_NRT is True - # Should restore to True (pre-context value), not to False. - assert numba_config.CUDA_ENABLE_NRT is True - finally: - numba_config.CUDA_ENABLE_NRT = saved - - -def test_nrt_enabled_restores_on_exception(): - """An exception inside `nrt_enabled()` must still restore the prior value.""" - import cudf.utils._numba # noqa: F401 - - assert numba_config.CUDA_ENABLE_NRT is False - with pytest.raises(RuntimeError, match="boom"): - with nrt_enabled(): - assert numba_config.CUDA_ENABLE_NRT is True - raise RuntimeError("boom") - assert numba_config.CUDA_ENABLE_NRT is False - - -def test_capture_nrt_usage_default_false(): - """`CaptureNRTUsage` starts with `use_nrt = False`.""" - cap = CaptureNRTUsage() - assert cap.use_nrt is False - - -def test_capture_nrt_usage_observes_inner_set(): - """Code inside the `with` block can flip `use_nrt` and the captor sees it.""" - with CaptureNRTUsage() as cap: - # A type instantiation that needs NRT would do this internally: - ctx = _current_nrt_context.get(None) - assert ctx is cap - ctx.use_nrt = True - assert cap.use_nrt is True +pytest.importorskip("numba_cuda_mlir") -def test_capture_nrt_usage_context_var_unset_outside(): - """The context var is reset on exit so nested code can't observe a stale captor.""" - with CaptureNRTUsage(): - pass - assert _current_nrt_context.get(None) is None +def test_mlir_backend_package_importable(): + import cudf.core.udf.mlir_backend as m -def test_capture_nrt_usage_nested(): - """Nesting two captors restores the outer context on inner exit.""" - with CaptureNRTUsage() as outer: - with CaptureNRTUsage() as inner: - assert _current_nrt_context.get(None) is inner - inner.use_nrt = True - assert _current_nrt_context.get(None) is outer - assert outer.use_nrt is False - assert inner.use_nrt is True + # Module is intentionally empty at this layer, only the docstring + # is present. + assert m.__doc__ diff --git a/python/cudf/cudf/utils/_numba.py b/python/cudf/cudf/utils/_numba.py index 8d62f8e7498f..79fdb1b1fc65 100644 --- a/python/cudf/cudf/utils/_numba.py +++ b/python/cudf/cudf/utils/_numba.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -6,13 +6,6 @@ from numba.cuda import config as numba_config from packaging import version -# Default NRT off for all numba-CUDA compilations driven through cudf. The -# UDF compile path (`cudf.core.udf.nrt_utils.nrt_enabled`) flips this to True -# per-kernel for UDFs whose data model declares an NRT meminfo (e.g. string -# returns), and restores it afterwards. Setting the global default to False -# avoids linking NRT into kernels that don't need it. -numba_config.CUDA_ENABLE_NRT = False - # Avoids using contextlib.contextmanager due to additional overhead class _CUDFNumbaConfig: