From 45c61ccd0cf9e16fd02c5bdd083b609e31738c17 Mon Sep 17 00:00:00 2001
From: Leo Fang <leof@nvidia.com>
Date: Fri, 6 Jun 2025 14:19:36 +0000
Subject: [PATCH] add PTX helpers

---
 cuda_bindings/cuda/bindings/__init__.py       |   1 +
 cuda_bindings/cuda/bindings/utils/__init__.py | 131 ++++++++++++++++++
 2 files changed, 132 insertions(+)
 create mode 100644 cuda_bindings/cuda/bindings/utils/__init__.py

diff --git a/cuda_bindings/cuda/bindings/__init__.py b/cuda_bindings/cuda/bindings/__init__.py
index 8086b4f6f..4f4f4032d 100644
--- a/cuda_bindings/cuda/bindings/__init__.py
+++ b/cuda_bindings/cuda/bindings/__init__.py
@@ -2,3 +2,4 @@
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 from cuda.bindings._version import __version__
+from cuda.bindings import utils
diff --git a/cuda_bindings/cuda/bindings/utils/__init__.py b/cuda_bindings/cuda/bindings/utils/__init__.py
new file mode 100644
index 000000000..adf643b77
--- /dev/null
+++ b/cuda_bindings/cuda/bindings/utils/__init__.py
@@ -0,0 +1,131 @@
+# Copyright 2025 NVIDIA Corporation.  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_driver_ver_from_ptx_ver(ptx_version: str) -> int:
+    """
+    Maps PTX ISA version to the minimal CUDA driver 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 driver 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}")
+
+
+
+# Regex pattern to match .version directive and capture the version number
+_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?")