NVTX Ranges in Numba Code Example

This example shows how to annotate Numba JIT-compiled code with NVTX ranges using the accvlab.optim_test_tools.numba_nvtx module. Numba-compiled functions run as native code; this module allows you to push and pop NVTX ranges inside such functions.

See also

The code of this example can be found in the repository under packages/optim_test_tools/examples/numba_nvtx_example.py.

Important

The functionality described here is specifically aimed at adding NVTX ranges to Numba-compiled code, where alternative solutions, e.g. the Python nvtx module or torch.cuda.nvtx, cannot be used. While accvlab.optim_test_tools.numba_nvtx also works for plain Python code, it is not intended for that use-case and does not provide benefits over using available alternatives in this case.

For profiling Python code, please also see the NVTX Range Wrapper Example example.

Overview

  • Import accvlab.optim_test_tools.numba_nvtx (e.g. as nvtx).

  • Register range names with register_string("...") to obtain integer handles; this must be done outside the JIT function, before compilation.

  • Inside an @numba.njit function, call nvtx.range_push(handle) and nvtx.range_pop() around the region to profile.

Example

Please see the notes in the code for more details.

packages/optim_test_tools/examples/numba_nvtx_example.py
 1# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 2#
 3# Licensed under the Apache License, Version 2.0 (the "License");
 4# you may not use this file except in compliance with the License.
 5# You may obtain a copy of the License at
 6#
 7#     http://www.apache.org/licenses/LICENSE-2.0
 8#
 9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import accvlab.optim_test_tools.numba_nvtx as nvtx
16from numba import njit
17
18# @NOTE
19# Register range names outside the JIT function to obtain integer handles.
20# Handles must be created before compilation because Numba compiles the function
21# and the handle value is baked into the compiled code.
22h_example_range = nvtx.register_string("example_range")
23h_example_range_inner = nvtx.register_string("example_range_inner")
24
25
26@njit
27def compute(x):
28    # @NOTE: Push the range using the handle created above.
29    nvtx.range_push(h_example_range)
30    y = x - 1
31    # @NOTE: Push the inner range.
32    nvtx.range_push(h_example_range_inner)
33    y = x + 2
34    # @NOTE: Pop both ranges.
35    nvtx.range_pop()
36    nvtx.range_pop()
37    return y
38
39
40if __name__ == "__main__":
41    result = compute(41)