Add PTX helpers

cuda_bindings/cuda/bindings/__init__.py

@@ -1,4 +1,5 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
+from cuda.bindings import utils
 from cuda.bindings._version import __version__

cuda_bindings/cuda/bindings/utils/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+from ._ptx_utils import get_minimal_required_cuda_ver_from_ptx_ver, get_ptx_ver

cuda_bindings/cuda/bindings/utils/_ptx_utils.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import re
+
+# Mapping based on the official PTX ISA <-> CUDA Release table
+# https://docs.nvidia.com/cuda/parallel-thread-execution/#release-notes-ptx-release-history
+_ptx_to_cuda = {
+    "1.0": (1, 0),
+    "1.1": (1, 1),
+    "1.2": (2, 0),
+    "1.3": (2, 1),
+    "1.4": (2, 2),
+    "2.0": (3, 0),
+    "2.1": (3, 1),
+    "2.2": (3, 2),
+    "2.3": (4, 0),
+    "3.0": (4, 1),
+    "3.1": (5, 0),
+    "3.2": (5, 5),
+    "4.0": (6, 0),
+    "4.1": (6, 5),
+    "4.2": (7, 0),
+    "4.3": (7, 5),
+    "5.0": (8, 0),
+    "6.0": (9, 0),
+    "6.1": (9, 1),
+    "6.2": (9, 2),
+    "6.3": (10, 0),
+    "6.4": (10, 1),
+    "6.5": (10, 2),
+    "7.0": (11, 0),
+    "7.1": (11, 1),
+    "7.2": (11, 2),
+    "7.3": (11, 3),
+    "7.4": (11, 4),
+    "7.5": (11, 5),
+    "7.6": (11, 6),
+    "7.7": (11, 7),
+    "7.8": (11, 8),
+    "8.0": (12, 0),
+    "8.1": (12, 1),
+    "8.2": (12, 2),
+    "8.3": (12, 3),
+    "8.4": (12, 4),
+    "8.5": (12, 5),
+    "8.6": (12, 7),
+    "8.7": (12, 8),
+    "8.8": (12, 9),
+}
+
+
+def get_minimal_required_cuda_ver_from_ptx_ver(ptx_version: str) -> int:
+    """
+    Maps the PTX ISA version to the minimal CUDA driver, nvPTXCompiler, or nvJitLink version
+    that is needed to load a PTX of the given ISA version.
+
+    Parameters
+    ----------
+    ptx_version : str
+        PTX ISA version as a string, e.g. "8.8" for PTX ISA 8.8. This is the ``.version``
+        directive in the PTX header.
+
+    Returns
+    -------
+    int
+        Minimal CUDA version as 1000 * major + 10 * minor, e.g. 12090 for CUDA 12.9.
+
+    Raises
+    ------
+    ValueError
+        If the PTX version is unknown.
+
+    Examples
+    --------
+    >>> get_minimal_required_driver_ver_from_ptx_ver("8.8")
+    12090
+    >>> get_minimal_required_driver_ver_from_ptx_ver("7.0")
+    11000
+    """
+    try:
+        major, minor = _ptx_to_cuda[ptx_version]
+        return 1000 * major + 10 * minor
+    except KeyError:
+        raise ValueError(f"Unknown or unsupported PTX ISA version: {ptx_version}") from None
+
+
+# Regex pattern to match .version directive and capture the version number
+# TODO: if import speed is a concern, consider lazy-initializing it.
+_ptx_ver_pattern = re.compile(r"\.version\s+([0-9]+\.[0-9]+)")
+
+
+def get_ptx_ver(ptx: str) -> str:
+    """
+    Extract the PTX ISA version string from PTX source code.
+
+    Parameters
+    ----------
+    ptx : str
+        The PTX assembly source code as a string.
+
+    Returns
+    -------
+    str
+        The PTX ISA version string, e.g., "8.8".
+
+    Raises
+    ------
+    ValueError
+        If the .version directive is not found in the PTX source.
+
+    Examples
+    --------
+    >>> ptx = r'''
+    ... .version 8.8
+    ... .target sm_86
+    ... .address_size 64
+    ...
+    ... .visible .entry test_kernel()
+    ... {
+    ...     ret;
+    ... }
+    ... '''
+    >>> get_ptx_ver(ptx)
+    '8.8'
+    """
+    m = _ptx_ver_pattern.search(ptx)
+    if m:
+        return m.group(1)
+    else:
+        raise ValueError("No .version directive found in PTX source. Is it a valid PTX?")

cuda_bindings/docs/source/api.rst

@@ -15,3 +15,4 @@ CUDA Python API Reference
    module/nvjitlink
    module/nvvm
    module/cufile
+   module/utils

cuda_bindings/docs/source/module/utils.rst

@@ -0,0 +1,16 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+.. module:: cuda.bindings.utils
+
+Utils module
+============
+
+Functions
+---------
+
+.. autosummary::
+   :toctree: generated/
+
+   get_minimal_required_cuda_ver_from_ptx_ver
+   get_ptx_ver

cuda_bindings/docs/source/release/12.X.Y-notes.rst

@@ -1,6 +1,8 @@
 .. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 .. SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
+.. module:: cuda.bindings
+
 ``cuda-bindings`` 12.X.Y Release notes
 ======================================
 
@@ -24,6 +26,8 @@ Bug fixes
 Miscellaneous
 -------------
 
+* Added PTX utilities including :func:`~utils.get_minimal_required_cuda_ver_from_ptx_ver` and :func:`~utils.get_ptx_ver`.
+
 
 Known issues
 ------------

cuda_bindings/tests/test_utils.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
+
+import pytest
+
+from cuda.bindings.utils import get_minimal_required_cuda_ver_from_ptx_ver, get_ptx_ver
+
+ptx_88_kernel = r"""
+.version 8.8
+.target sm_75
+.address_size 64
+
+	// .globl	empty_kernel
+
+.visible .entry empty_kernel()
+{
+	ret;
+}
+"""
+
+
+ptx_72_kernel = r"""
+.version  7.2
+.target sm_75
+.address_size 64
+
+	// .globl	empty_kernel
+
+.visible .entry empty_kernel()
+{
+	ret;
+}
+"""
+
+
+@pytest.mark.parametrize(
+    "kernel,actual_ptx_ver,min_cuda_ver", ((ptx_88_kernel, "8.8", 12090), (ptx_72_kernel, "7.2", 11020))
+)
+def test_ptx_utils(kernel, actual_ptx_ver, min_cuda_ver):
+    ptx_ver = get_ptx_ver(kernel)
+    assert ptx_ver == actual_ptx_ver
+    cuda_ver = get_minimal_required_cuda_ver_from_ptx_ver(ptx_ver)
+    assert cuda_ver == min_cuda_ver