feat(lower-tirx): align NVIDIA IKET profiling with the official ABI

.gitignore

@@ -183,6 +183,7 @@ cscope*
 perf
 .bash_history
 *.json
+!tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json
 *.params
 *.ro
 *.onnx

python/tvm/backend/cuda/iket.py

@@ -52,24 +52,26 @@
 from tvm.script import tirx as T
 
 _PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE"
-_DEFAULT_PROFILE = "cutlass-4.6.1"
+_INJECTED_CHILD_ENABLE_ENV = "TVM_IKET_INJECTED_CHILD_ENABLE"
+_DEFAULT_PROFILE = "cutlass-4.6.0"
 _POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all"))
 _INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG")
 _OUTPUT_TAIL_LINES = 100
 _TERMINATION_GRACE_SECONDS = 5.0
 
-# Hashes are SHA-256 digests of files in the public 4.6.1 wheels.  Only
+# Hashes are SHA-256 digests of files in the public 4.6.0 wheels.  Only
 # ABI-independent runtime/compiler binaries are pinned so this profile works
 # with every Python version supported by that CUTLASS DSL release.
 _OFFICIAL_PROFILES = {
-    "cutlass-4.6.1": {
+    "cutlass-4.6.0": {
+        "nvrtc_version": (13, 2),
         "versions": {
-            "nvidia-cutlass-dsl": "4.6.1",
-            "nvidia-cutlass-dsl-libs-base": "4.6.1",
-            "nvidia-cutlass-dsl-libs-core": "4.6.1",
-            "nvidia-cutlass-dsl-libs-cu13": "4.6.1",
+            "nvidia-cutlass-dsl": "4.6.0",
+            "nvidia-cutlass-dsl-libs-base": "4.6.0",
+            "nvidia-cutlass-dsl-libs-core": "4.6.0",
+            "nvidia-cutlass-dsl-libs-cu13": "4.6.0",
             "nvidia-cuda-nvdisasm": "13.3.73",
-            "nvidia-cuda-nvrtc": "13.3.33",
+            "nvidia-cuda-nvrtc": "13.2.78",
         },
         "files": {
             "nvidia-cutlass-dsl-libs-base": {
@@ -92,10 +94,10 @@
             },
             "nvidia-cuda-nvrtc": {
                 "nvidia/cu13/lib/libnvrtc.so.13": (
-                    "e51d197b3b0d2d9d850d29977423e6ac60661d429a59c440fc04e52b6fc6750a"
+                    "c673cf3b5099d83b98a388a2bb21e5d6f481be3c4bb956e2d74c39cb714d8c63"
                 ),
-                "nvidia/cu13/lib/libnvrtc-builtins.so.13.3": (
-                    "7394c640e5761d13d2bbcdbc4b4c5dbac7cb53cd5bc732d78f8a5cb38638e913"
+                "nvidia/cu13/lib/libnvrtc-builtins.so.13.2": (
+                    "6b1c571cc730d5fcfd57f322e1fa7e0e65de7454b2239ff6d552a09b82d47dbe"
                 ),
             },
         },
@@ -218,19 +220,23 @@ def _validate_run_iket_entrypoint() -> str:
         and item.value == "iket.cli.main:entrypoint"
         for item in entry_points
     ):
-        raise _profile_error("the run-iket entry point does not match CUTLASS DSL 4.6.1")
+        raise _profile_error("the run-iket entry point does not match the locked CUTLASS profile")
     return executable
 
 
-def _validate_nvrtc_13_3() -> None:
+def _validate_nvrtc_version(expected_version: tuple[int, int]) -> None:
+    expected_label = ".".join(str(part) for part in expected_version)
     try:
         from cuda.bindings import nvrtc
 
         error, major, minor = nvrtc.nvrtcVersion()
     except (ImportError, OSError, RuntimeError) as err:
-        raise _profile_error("CUDA NVRTC 13.3 is unavailable") from err
-    if int(error) != 0 or (int(major), int(minor)) != (13, 3):
-        raise _profile_error(f"CUDA NVRTC 13.3 is required, got {int(major)}.{int(minor)}")
+        raise _profile_error(f"CUDA NVRTC {expected_label} is unavailable") from err
+    actual_version = (int(major), int(minor))
+    if int(error) != 0 or actual_version != expected_version:
+        raise _profile_error(
+            f"CUDA NVRTC {expected_label} is required, got {actual_version[0]}.{actual_version[1]}"
+        )
 
 
 def _validate_official_installation(profile_name: str) -> str:
@@ -265,7 +271,7 @@ def _validate_official_installation(profile_name: str) -> str:
                 )
 
     executable = _validate_run_iket_entrypoint()
-    _validate_nvrtc_13_3()
+    _validate_nvrtc_version(profile_config["nvrtc_version"])
     return executable
 
 
@@ -363,17 +369,28 @@ def __call__(self, *args, **kwargs):
 class IketProfiler:
     """TIRx annotations compiled for NVIDIA's official IKET runtime."""
 
-    def mark(self, name: str):
-        T.evaluate(T.cuda.iket.mark(name))
-
-    def range_start(self, name: str):
-        return T.cuda.iket.range_start(name)
-
-    def range_end(self, token: tvm.tirx.Expr):
-        T.evaluate(T.cuda.iket.range_end(token))
-
-    def range_push(self, name: str):
-        T.evaluate(T.cuda.iket.range_push(name))
+    def mark(self, name: str, payload=None):
+        if payload is None:
+            T.evaluate(T.cuda.iket.mark(name))
+        else:
+            T.evaluate(T.cuda.iket.mark(name, payload))
+
+    def range_start(self, name: str, payload=None):
+        if payload is None:
+            return T.cuda.iket.range_start(name)
+        return T.cuda.iket.range_start(name, payload)
+
+    def range_end(self, token: tvm.tirx.Expr, payload=None):
+        if payload is None:
+            T.evaluate(T.cuda.iket.range_end(token))
+        else:
+            T.evaluate(T.cuda.iket.range_end(token, payload))
+
+    def range_push(self, name: str, payload=None):
+        if payload is None:
+            T.evaluate(T.cuda.iket.range_push(name))
+        else:
+            T.evaluate(T.cuda.iket.range_push(name, payload))
 
     def range_pop(self):
         T.evaluate(T.cuda.iket.range_pop())
@@ -450,6 +467,10 @@ def _child_environment(
             stacklevel=3,
         )
     child_env[_PROFILE_ENV] = profile_name
+    # LowerIket also requires the two run-iket injection variables before it
+    # honors this marker.  This enables ordinary TIRx JIT compilation only in
+    # children started by this locked profiling entry point.
+    child_env[_INJECTED_CHILD_ENABLE_ENV] = "1"
     return child_env
 
 

python/tvm/backend/cuda/intrinsics/misc.py

@@ -197,14 +197,15 @@ def _write_event(event_bits: str) -> str:
 # Official IKET NativeDump placeholder.
 # =============================================================================
 @register_codegen("cuda_iket_official_event")
-def codegen_cuda_iket_official_event(event_id, source_code):
+def codegen_cuda_iket_official_event(event_id, source_code, payload=None):
     if isinstance(source_code, tvm.tirx.StringImm):
         source_code = source_code.value
     else:
         source_code = parse_str(source_code)
+    args = (event_id,) if payload is None else (event_id, payload)
     return cuda_func_call(
         "tvm_builtin_iket_official_event",
-        event_id,
+        *args,
         source_code=source_code,
         return_type="uint32",
     )

python/tvm/backend/cuda/op.py

@@ -65,23 +65,31 @@
 ########################################################
 
 
-def cuda_iket_mark(name):
+def cuda_iket_mark(name, payload=None):
     """Create an NVIDIA IKET marker annotation."""
+    if payload is not None:
+        return call_intrin("", "tirx.cuda.iket_mark", name, payload)
     return call_intrin("", "tirx.cuda.iket_mark", name)
 
 
-def cuda_iket_range_start(name):
+def cuda_iket_range_start(name, payload=None):
     """Create an NVIDIA IKET token-range start annotation."""
+    if payload is not None:
+        return call_intrin("uint32", "tirx.cuda.iket_range_start", name, payload)
     return call_intrin("uint32", "tirx.cuda.iket_range_start", name)
 
 
-def cuda_iket_range_end(token):
+def cuda_iket_range_end(token, payload=None):
     """Create an NVIDIA IKET token-range end annotation."""
+    if payload is not None:
+        return call_intrin("", "tirx.cuda.iket_range_end", token, payload)
     return call_intrin("", "tirx.cuda.iket_range_end", token)
 
 
-def cuda_iket_range_push(name):
+def cuda_iket_range_push(name, payload=None):
     """Create an NVIDIA IKET stack-range push annotation."""
+    if payload is not None:
+        return call_intrin("", "tirx.cuda.iket_range_push", name, payload)
     return call_intrin("", "tirx.cuda.iket_range_push", name)
 
 
@@ -95,8 +103,12 @@ def cuda_iket_sentinel_token(name):
     return call_intrin("uint32", "tirx.cuda.iket_sentinel_token", name)
 
 
-def cuda_iket_official_event(event_id, source_code=""):
+def cuda_iket_official_event(event_id, source_code="", payload=None):
     """Create an NVIDIA IKET official range-end event."""
+    if payload is not None:
+        return call_intrin(
+            "uint32", "tirx.cuda.iket_official_event", event_id, source_code, payload
+        )
     return call_intrin("uint32", "tirx.cuda.iket_official_event", event_id, source_code)