NVIDIA · gmarkall · Oct 14, 2025 · Oct 3, 2025 · Oct 3, 2025 · Oct 9, 2025
diff --git a/numba_cuda/numba/cuda/core/annotations/pretty_annotate.py b/numba_cuda/numba/cuda/core/annotations/pretty_annotate.py
@@ -0,0 +1,288 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+This module implements code highlighting of numba-cuda function annotations.
+"""
+
+from warnings import warn
+
+warn(
+    "The pretty_annotate functionality is experimental and might change API",
+    FutureWarning,
+)
+
+
+def hllines(code, style):
+    try:
+        from pygments import highlight
+        from pygments.lexers import PythonLexer
+        from pygments.formatters import HtmlFormatter
+    except ImportError:
+        raise ImportError("please install the 'pygments' package")
+    pylex = PythonLexer()
+    "Given a code string, return a list of html-highlighted lines"
+    hf = HtmlFormatter(noclasses=True, style=style, nowrap=True)
+    res = highlight(code, pylex, hf)
+    return res.splitlines()
+
+
+def htlines(code, style):
+    try:
+        from pygments import highlight
+        from pygments.lexers import PythonLexer
+
+        # TerminalFormatter does not support themes, Terminal256 should,
+        # but seem to not work.
+        from pygments.formatters import TerminalFormatter
+    except ImportError:
+        raise ImportError("please install the 'pygments' package")
+    pylex = PythonLexer()
+    "Given a code string, return a list of ANSI-highlighted lines"
+    hf = TerminalFormatter(style=style)
+    res = highlight(code, pylex, hf)
+    return res.splitlines()
+
+
+def get_ansi_template():
+    try:
+        from jinja2 import Template
+    except ImportError:
+        raise ImportError("please install the 'jinja2' package")
+    return Template("""
+    {%- for func_key in func_data.keys() -%}
+        Function name: \x1b[34m{{func_data[func_key]['funcname']}}\x1b[39;49;00m
+        {%- if func_data[func_key]['filename'] -%}
+        {{'\n'}}In file: \x1b[34m{{func_data[func_key]['filename'] -}}\x1b[39;49;00m
+        {%- endif -%}
+        {{'\n'}}With signature: \x1b[34m{{func_key[1]}}\x1b[39;49;00m
+        {{- "\n" -}}
+        {%- for num, line, hl, hc in func_data[func_key]['pygments_lines'] -%}
+                {{-'\n'}}{{ num}}: {{hc-}}
+                {%- if func_data[func_key]['ir_lines'][num] -%}
+                    {%- for ir_line, ir_line_type in func_data[func_key]['ir_lines'][num] %}
+                        {{-'\n'}}--{{- ' '*func_data[func_key]['python_indent'][num]}}
+                        {{- ' '*(func_data[func_key]['ir_indent'][num][loop.index0]+4)
+                        }}{{ir_line }}\x1b[41m{{ir_line_type-}}\x1b[39;49;00m
+                    {%- endfor -%}
+                {%- endif -%}
+            {%- endfor -%}
+    {%- endfor -%}
+    """)
+
+
+def get_html_template():
+    try:
+        from jinja2 import Template
+    except ImportError:
+        raise ImportError("please install the 'jinja2' package")
+    return Template("""
+    <html>
+    <head>
+        <style>
+
+            .annotation_table {
+                color: #000000;
+                font-family: monospace;
+                margin: 5px;
+                width: 100%;
+            }
+
+            /* override JupyterLab style */
+            .annotation_table td {
+                text-align: left;
+                background-color: transparent;
+                padding: 1px;
+            }
+
+            .annotation_table tbody tr:nth-child(even) {
+                background: white;
+            }
+
+            .annotation_table code
+            {
+                background-color: transparent;
+                white-space: normal;
+            }
+
+            /* End override JupyterLab style */
+
+            tr:hover {
+                background-color: rgba(92, 200, 249, 0.25);
+            }
+
+            td.object_tag summary ,
+            td.lifted_tag summary{
+                font-weight: bold;
+                display: list-item;
+            }
+
+            span.lifted_tag {
+                color: #00cc33;
+            }
+
+            span.object_tag {
+                color: #cc3300;
+            }
+
+
+            td.lifted_tag {
+                background-color: #cdf7d8;
+            }
+
+            td.object_tag {
+                background-color: #fef5c8;
+            }
+
+            code.ir_code {
+                color: grey;
+                font-style: italic;
+            }
+
+            .metadata {
+                border-bottom: medium solid black;
+                display: inline-block;
+                padding: 5px;
+                width: 100%;
+            }
+
+            .annotations {
+                padding: 5px;
+            }
+
+            .hidden {
+                display: none;
+            }
+
+            .buttons {
+                padding: 10px;
+                cursor: pointer;
+            }
+        </style>
+    </head>
+
+    <body>
+        {% for func_key in func_data.keys() %}
+            <div class="metadata">
+            Function name: {{func_data[func_key]['funcname']}}<br />
+            {% if func_data[func_key]['filename'] %}
+                in file: {{func_data[func_key]['filename']|escape}}<br />
+            {% endif %}
+            with signature: {{func_key[1]|e}}
+            </div>
+            <div class="annotations">
+            <table class="annotation_table tex2jax_ignore">
+                {%- for num, line, hl, hc in func_data[func_key]['pygments_lines'] -%}
+                    {%- if func_data[func_key]['ir_lines'][num] %}
+                        <tr><td style="text-align:left;" class="{{func_data[func_key]['python_tags'][num]}}">
+                            <details>
+                                <summary>
+                                    <code>
+                                    {{num}}:
+                                    {{'&nbsp;'*func_data[func_key]['python_indent'][num]}}{{hl}}
+                                    </code>
+                                </summary>
+                                <table class="annotation_table">
+                                    <tbody>
+                                        {%- for ir_line, ir_line_type in func_data[func_key]['ir_lines'][num] %}
+                                            <tr class="ir_code">
+                                                <td style="text-align: left;"><code>
+                                                &nbsp;
+                                                {{- '&nbsp;'*func_data[func_key]['python_indent'][num]}}
+                                                {{ '&nbsp;'*func_data[func_key]['ir_indent'][num][loop.index0]}}{{ir_line|e -}}
+                                                <span class="object_tag">{{ir_line_type}}</span>
+                                                </code>
+                                                </td>
+                                            </tr>
+                                        {%- endfor -%}
+                                    </tbody>
+                                </table>
+                                </details>
+                        </td></tr>
+                    {% else -%}
+                        <tr><td style="text-align:left; padding-left: 22px;" class="{{func_data[func_key]['python_tags'][num]}}">
+                            <code>
+                                {{num}}:
+                                {{'&nbsp;'*func_data[func_key]['python_indent'][num]}}{{hl}}
+                            </code>
+                        </td></tr>
+                    {%- endif -%}
+                {%- endfor -%}
+            </table>
+            </div>
+        {% endfor %}
+    </body>
+    </html>
+    """)
+
+
+def reform_code(annotation):
+    """
+    Extract the code from the Numba-cuda annotation datastructure.
+
+    Pygments can only highlight full multi-line strings, the Numba-cuda
+    annotation is list of single lines, with indentation removed.
+    """
+    ident_dict = annotation["python_indent"]
+    s = ""
+    for n, l in annotation["python_lines"]:
+        s = s + " " * ident_dict[n] + l + "\n"
+    return s
+
+
+class Annotate:
+    """
+    Construct syntax highlighted annotation for a given jitted function:
+
+    Example:
+
+    >>> from numba import cuda
+    >>> import numpy as np
+    >>> from numba.cuda.core.annotations.pretty_annotate import Annotate
+    >>> @cuda.jit
+    ... def test(a):
+    ...     tid = cuda.grid(1)
+    ...     size = len(a)
+    ...     if tid < size:
+    ...         a[tid] = 1
+    >>> test[(4), (16)](np.ones(100))
+    >>> Annotate(test)
+
+    The last line will return an HTML and/or ANSI representation that will be
+    displayed accordingly in Jupyter/IPython.
+
+    Function annotations persist across compilation for newly encountered
+    type signatures and as a result annotations are shown for all signatures
+    by default.
+
+    Annotations for a specific signature can be shown by using the
+    ``signature`` parameter. For the above jitted function:
+
+    >>> test.signatures
+    [(Array(float64, 1, 'C', False, aligned=True),)]
+    >>> Annotate(f, signature=f.signatures[0])
+    # annotation for Array(float64, 1, 'C', False, aligned=True)
+    """
+
+    def __init__(self, function, signature=None, **kwargs):
+        style = kwargs.get("style", "default")
+        if not function.signatures:
+            raise ValueError(
+                "function need to be jitted for at least one signature"
+            )
+        ann = function.get_annotation_info(signature=signature)
+        self.ann = ann
+
+        for k, v in ann.items():
+            res = hllines(reform_code(v), style)
+            rest = htlines(reform_code(v), style)
+            v["pygments_lines"] = [
+                (a, b, c, d)
+                for (a, b), c, d in zip(v["python_lines"], res, rest)
+            ]
+
+    def _repr_html_(self):
+        return get_html_template().render(func_data=self.ann)
+
+    def __repr__(self):
+        return get_ansi_template().render(func_data=self.ann)
diff --git a/numba_cuda/numba/cuda/core/annotations/type_annotations.py b/numba_cuda/numba/cuda/core/annotations/type_annotations.py
@@ -11,7 +11,6 @@
 import textwrap
 from io import StringIO
 
-import numba.core.dispatcher
 from numba.core import ir
 
 
@@ -83,6 +82,8 @@ def __init__(
         self.lifted_from = lifted_from
 
     def prepare_annotations(self):
+        from numba.cuda.dispatcher import LiftedLoop
+
         # Prepare annotations
         groupedinst = defaultdict(list)
         found_lifted_loop = False
@@ -103,7 +104,7 @@ def prepare_annotations(self):
                     ):
                         atype = self.calltypes[inst.value]
                     elif isinstance(inst.value, ir.Const) and isinstance(
-                        inst.value.value, numba.core.dispatcher.LiftedLoop
+                        inst.value.value, LiftedLoop
                     ):
                         atype = "XXX Lifted Loop XXX"
                         found_lifted_loop = True

diff --git a/numba_cuda/numba/cuda/core/entrypoints.py b/numba_cuda/numba/cuda/core/entrypoints.py
@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import logging
+import warnings
+
+from importlib import metadata as importlib_metadata
+
+
+_already_initialized = False
+logger = logging.getLogger(__name__)
+
+
+def init_all():
+    """Execute all `numba_cuda_extensions` entry points with the name `init`
+
+    If extensions have already been initialized, this function does nothing.
+    """
+    try:
+        from numba.core import entrypoints
+
+        entrypoints.init_all()
+    except ImportError:
+        pass
+
+    global _already_initialized
+    if _already_initialized:
+        return
+
+    # Must put this here to avoid extensions re-triggering initialization
+    _already_initialized = True
+
+    def load_ep(entry_point):
+        """Loads a given entry point. Warns and logs on failure."""
+        logger.debug("Loading extension: %s", entry_point)
+        try:
+            func = entry_point.load()
+            func()
+        except Exception as e:
+            msg = (
+                f"Numba extension module '{entry_point.module}' "
+                f"failed to load due to '{type(e).__name__}({str(e)})'."
+            )
+            warnings.warn(msg, stacklevel=3)
+            logger.debug("Extension loading failed for: %s", entry_point)
+
+    eps = importlib_metadata.entry_points()
+    # Split, Python 3.10+ and importlib_metadata 3.6+ have the "selectable"
+    # interface, versions prior to that do not. See "compatibility note" in:
+    # https://docs.python.org/3.10/library/importlib.metadata.html#entry-points
+    if hasattr(eps, "select"):
+        for entry_point in eps.select(
+            group="numba_cuda_extensions", name="init"
+        ):
+            load_ep(entry_point)
+    else:
+        for entry_point in eps.get("numba_cuda_extensions", ()):
+            if entry_point.name == "init":
+                load_ep(entry_point)
diff --git a/numba_cuda/numba/cuda/core/transforms.py b/numba_cuda/numba/cuda/core/transforms.py
@@ -193,7 +193,7 @@ def _loop_lift_modify_blocks(
     Modify the block inplace to call to the lifted-loop.
     Returns a dictionary of blocks of the lifted-loop.
     """
-    from numba.core.dispatcher import LiftedLoop
+    from numba.cuda.dispatcher import LiftedLoop
 
     # Copy loop blocks
     loop = loopinfo.loop
@@ -402,7 +402,7 @@ def with_lifting(func_ir, typingctx, targetctx, flags, locals):
     from numba.cuda.core import postproc
 
     def dispatcher_factory(func_ir, objectmode=False, **kwargs):
-        from numba.core.dispatcher import LiftedWith, ObjModeLiftedWith
+        from numba.cuda.dispatcher import LiftedWith, ObjModeLiftedWith
 
         myflags = flags.copy()
         if objectmode: