fixup! feat(core.utils): canonicalize ptxas_options; scope use_libdevice to NVVM

Author: cpcloud

cuda_core/cuda/core/utils/_program_cache.py

@@ -241,6 +241,19 @@ def _gate_identity(v):
     return v
 
 
+def _gate_ptxas_options(v):
+    # ``_prepare_nvjitlink_options`` emits one ``-Xptxas=<s>`` per element, and
+    # treats ``str`` as a single-element sequence. Canonicalize to a tuple so
+    # ``"-v"`` / ``["-v"]`` / ``("-v",)`` all hash the same.
+    if v is None:
+        return None
+    if isinstance(v, str):
+        return ("-Xptxas=" + v,)
+    if isinstance(v, collections.abc.Sequence):
+        return tuple(f"-Xptxas={s}" for s in v)
+    return v
+
+
 _LINKER_FIELD_GATES = {
     "name": _gate_identity,
     "arch": _gate_identity,
@@ -254,7 +267,7 @@ def _gate_identity(v):
     "prec_sqrt": _gate_tristate_bool,
     "fma": _gate_tristate_bool,
     "split_compile": _gate_identity,
-    "ptxas_options": _gate_identity,
+    "ptxas_options": _gate_ptxas_options,
     "no_cache": _gate_is_true,
 }
 
@@ -655,10 +668,11 @@ def _probe(label: str, fn):
             else:
                 # Fallback for unexpected format.
                 _update("extra_source", str(item).encode("utf-8"))
-    # Program_init gates ``use_libdevice`` on truthiness, not ``is not None``
-    # (see _program.pyx), so False and None produce identical ObjectCode and
-    # must hash the same way here.
-    if getattr(options, "use_libdevice", None):
+    # ``use_libdevice`` is only consumed on the NVVM compile path
+    # (_program.pyx Program_init); NVRTC and PTX/linker ignore it, so
+    # folding it into the key there would force spurious misses. On NVVM,
+    # Program_init gates it on truthiness -- False and None match.
+    if backend == "nvvm" and getattr(options, "use_libdevice", None):
         _update("use_libdevice", b"1")
 
     # Program.compile() propagates options.name onto the returned ObjectCode,

cuda_core/tests/test_program_cache.py

@@ -395,6 +395,26 @@ def test_make_program_cache_key_ptx_driver_ignored_fields_collapse(field, a, b,
     assert k_a == k_b
 
 
+@pytest.mark.parametrize(
+    "a, b",
+    [
+        pytest.param("-v", ["-v"], id="str_vs_list"),
+        pytest.param("-v", ("-v",), id="str_vs_tuple"),
+        pytest.param(["-v"], ("-v",), id="list_vs_tuple"),
+    ],
+)
+def test_make_program_cache_key_ptx_ptxas_options_canonicalized(a, b, monkeypatch):
+    """_prepare_nvjitlink_options emits the same -Xptxas= flags for str,
+    list, and tuple shapes of ptxas_options. The cache key must treat them
+    as equivalent so equivalent compiles don't miss the cache."""
+    from cuda.core import _linker
+
+    monkeypatch.setattr(_linker, "_decide_nvjitlink_or_driver", lambda: False)  # nvJitLink
+    k_a = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ptxas_options=a))
+    k_b = _make_key(code=".version 7.0", code_type="ptx", options=_opts(ptxas_options=b))
+    assert k_a == k_b
+
+
 def test_make_program_cache_key_ptx_driver_ignored_fields_still_matter_under_nvjitlink(monkeypatch):
     """nvJitLink does honour those fields; they must still differentiate keys there."""
     from cuda.core import _linker
@@ -405,6 +425,22 @@ def test_make_program_cache_key_ptx_driver_ignored_fields_still_matter_under_nvj
     assert k_a != k_b
 
 
+@pytest.mark.parametrize(
+    "code_type, code, target_type",
+    [
+        pytest.param("c++", "void k(){}", "cubin", id="nvrtc"),
+        pytest.param("ptx", ".version 7.0", "cubin", id="ptx"),
+    ],
+)
+def test_make_program_cache_key_use_libdevice_ignored_for_non_nvvm(code_type, code, target_type):
+    """``use_libdevice`` is only consumed on the NVVM path; NVRTC and PTX
+    ignore it, so toggling it must not perturb the cache key elsewhere."""
+    k_off = _make_key(code=code, code_type=code_type, target_type=target_type, options=_opts(use_libdevice=False))
+    k_on = _make_key(code=code, code_type=code_type, target_type=target_type, options=_opts(use_libdevice=True))
+    k_none = _make_key(code=code, code_type=code_type, target_type=target_type, options=_opts(use_libdevice=None))
+    assert k_off == k_on == k_none
+
+
 def test_make_program_cache_key_nvvm_use_libdevice_false_equals_none():
     """Program_init gates ``use_libdevice`` on truthiness, so False and None
     compile identically and must hash the same way."""