Skip to content

feat: support decomposing, indexing, and iterating Dim3 objects (fixes #134, #55) - #933

Open
reginaldalfret wants to merge 1 commit into
NVIDIA:mainfrom
reginaldalfret:fix/issue-134-decompose-dim3
Open

reginaldalfret wants to merge 1 commit into
NVIDIA:mainfrom
reginaldalfret:fix/issue-134-decompose-dim3

Conversation

@reginaldalfret

Copy link
Copy Markdown

Summary

In CUDA kernels and simulator code, thread and block dimensions/indices (cuda.threadIdx, cuda.blockDim, cuda.blockIdx, cuda.gridDim) return Dim3 objects. Previously, Dim3 was modeled as an opaque Type in the typing system with only .x, .y, .z attributes. Unpacking a Dim3 object (e.g. x, y, z = cuda.threadIdx or bx, by, bz = cuda.blockDim) failed with TypingError: failed to unpack Dim3. Furthermore, subscripting cuda.threadIdx[0] or using len() / for v in cuda.threadIdx: was unsupported. In the simulator, Dim3 had an __iter__ method but lacked __getitem__ and __len__, causing TypeError: 'Dim3' object is not subscriptable.

Root Cause

  • In numba_cuda/numba/cuda/types/ext_types.py, Dim3 inherited directly from Type rather than NamedUniTuple, lacking sequence/unpacking typing protocol semantics.
  • In numba_cuda/numba/cuda/models.py and cudaimpl.py, Dim3Model used StructModel and packed via cgutils.pack_struct rather than cgutils.pack_array matching UniTupleModel's contiguous homogeneous layout.
  • In numba_cuda/numba/cuda/simulator/kernelapi.py and stubs.py, Dim3 lacked __getitem__ and __len__ definitions.

Changes

  • Typing (ext_types.py): Defined _Dim3Tuple = namedtuple("Dim3", ("x", "y", "z")) and updated Dim3 to inherit from NamedUniTuple(Integer.from_bitwidth(32, signed=True), 3, _Dim3Tuple).
  • Data Model (models.py): Updated Dim3Model to subclass models.UniTupleModel.
  • Lowering (cudaimpl.py): Updated initialize_dim3 to pack elements using cgutils.pack_array(builder, (x, y, z)) matching UniTupleModel.
  • Simulator & Stubs (simulator/kernelapi.py, stubs.py): Added __getitem__ and __len__ methods to Dim3 for full feature parity across device and simulator targets.
  • Tests:
    • Added TestDim3Decompose in numba_cuda/numba/cuda/tests/cudapy/test_intrinsics.py testing test_decompose_threadIdx, test_decompose_blockDim, test_decompose_gridDim, test_indexing, test_len, test_iteration, and test_attributes_still_work.
    • Added test_dim3_sequence_support in numba_cuda/numba/cuda/tests/cudasim/test_cudasim_issues.py.

Tests

  • numba.cuda.tests.cudasim.test_cudasim_issues.TestCudaSimIssues: 5/5 passed
  • numba.cuda.tests.cudapy.test_intrinsics.TestDim3Decompose: 7/7 passed
  • numba.cuda.tests.cudapy.test_print.TestPrint.test_dim3: 1/1 passed
  • Direct LLVM IR compilation test (compile_cuda) verifying correct PTX register read and unpack code generation: PASSED
  • ruff check numba_cuda: 0 errors
  • git diff --check: clean

Validation

  • Tested locally in simulator mode and verified LLVM IR emission through compile_cuda.
  • GPU hardware execution is left to NVIDIA CI.

Issue

Fixes #134
Fixes #55

@copy-pr-bot

copy-pr-bot Bot commented Sep 23, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

This branch has not been deployed

No deployments
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.

[FEA] Decompose Dim3s (x, y, z = cuda.blockDim) [FEA] Dim3 should act more like a NamedTuple

1 participant