From 4505a6187f971b0955b64612030bcc7931290df4 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Mon, 3 Aug 2026 14:32:29 -0400 Subject: [PATCH 1/5] feat(lower-tirx): align official IKET profiling Lock the official integration to CUTLASS DSL 4.6.0 with NVRTC 13.2 and refresh its oracle. Report unproven warp convergence as an advisory warning and handle nested loop control independently. --- .gitignore | 1 + python/tvm/backend/cuda/iket.py | 39 +++++++----- src/backend/cuda/transforms/lower_iket.cc | 19 +++--- .../oracle/generate_iket_official_oracle.py | 2 +- ...> iket_official_cutlass_4_6_0_oracle.json} | 14 ++--- .../tirx/iket/test_iket_orchestration.py | 4 +- tests/python/tirx/iket/test_iket_profiler.py | 63 +++++++++++++------ 7 files changed, 89 insertions(+), 53 deletions(-) rename tests/python/tirx/iket/oracle/{iket_official_cutlass_4_6_1_oracle.json => iket_official_cutlass_4_6_0_oracle.json} (85%) diff --git a/.gitignore b/.gitignore index 2a59857ea7b0..b2f8ee1d984d 100644 --- a/.gitignore +++ b/.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 diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py index 1ccbfc11001b..c21d56835184 100644 --- a/python/tvm/backend/cuda/iket.py +++ b/python/tvm/backend/cuda/iket.py @@ -52,24 +52,25 @@ from tvm.script import tirx as T _PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE" -_DEFAULT_PROFILE = "cutlass-4.6.1" +_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 +93,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 +219,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 +270,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 diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc index 405174435a36..ccb7e9db2435 100644 --- a/src/backend/cuda/transforms/lower_iket.cc +++ b/src/backend/cuda/transforms/lower_iket.cc @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -1213,6 +1214,11 @@ class LoopControlFinder : public StmtExprVisitor { bool found{false}; private: + // A break or continue nested inside another loop targets that inner loop, + // not the loop whose body this finder was asked to inspect. Each nested + // loop is checked independently by IketConvergenceVerifier. + void VisitStmt_(const ForNode*) final {} + void VisitStmt_(const WhileNode*) final {} void VisitStmt_(const BreakNode* op) final { found = true; } void VisitStmt_(const ContinueNode* op) final { found = true; } void VisitExpr_(const CallNode* call) final { @@ -1382,13 +1388,11 @@ class IketConvergenceVerifier : public StmtExprVisitor { void VisitExpr_(const CallNode* call) final { if (IsIketOp(call->op)) { - if (divergent_context_) { - TVM_FFI_THROW(ValueError) << "IKET event site may be reached by a divergent set of lanes: " - << GetRef(call); - } - if (call->op.same_as(IketRangeEndOp())) { - TVM_FFI_CHECK(IsUniform(call->args[0]), ValueError) - << "IKET range_end requires a warp-uniform RangeToken"; + bool unproven_token = call->op.same_as(IketRangeEndOp()) && !IsUniform(call->args[0]); + if ((divergent_context_ || unproven_token) && warned_calls_.insert(call).second) { + LOG(WARNING) << "IKET warp convergence could not be proven for event site: " + << GetRef(call) + << "; continuing because convergence diagnostics are advisory"; } } StmtExprVisitor::VisitExpr_(call); @@ -1396,6 +1400,7 @@ class IketConvergenceVerifier : public StmtExprVisitor { DivergentVarSet divergent_vars_; UniformBufferSet uniform_buffers_; + std::unordered_set warned_calls_; bool divergent_context_{false}; }; diff --git a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py index acc5826381f1..2b4a99256477 100644 --- a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py +++ b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py @@ -13,7 +13,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Generate the normalized CUTLASS DSL 4.6.1 IKET oracle manifest. +"""Generate the normalized CUTLASS DSL 4.6.0 IKET oracle manifest. This helper intentionally writes NVIDIA-generated binary artifacts only to the requested output directory. The repository stores the normalized manifest, diff --git a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json similarity index 85% rename from tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json rename to tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json index d84d38115d73..0bdb9e50f142 100644 --- a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json +++ b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json @@ -190,18 +190,18 @@ "--keep-cubin", "--keep-sass" ], - "cutlass_dsl": "4.6.1", + "cutlass_dsl": "4.6.0", "driver": "NVIDIA B200, 595.58.03", "instrument_method": "NativeDump", - "nvdisasm": "nvdisasm: NVIDIA (R) CUDA disassembler\nCopyright (c) 2005-2026 NVIDIA Corporation\nBuilt on Tue_Jun_09_02:42:28_PM_PDT_2026\nCuda compilation tools, release 13.3, V13.3.73\nBuild cuda_13.3.r13.3/compiler.38244171_0" + "nvdisasm": "nvdisasm: NVIDIA (R) CUDA disassembler\nCopyright (c) 2005-2026 NVIDIA Corporation\nBuilt on Mon_Mar_02_09:52:52_PM_PST_2026\nCuda compilation tools, release 13.2, V13.2.51\nBuild cuda_13.2.r13.2/compiler.37434383_0" }, "schema_version": 2, "wheels": { "nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl": "fa17084b07c0dca68a42892f771b4b1b40fbe9b91660209623e61cea611cae8c", - "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl": "82530788b8c6164a54d3fd9ae8bcca8893d397c4aeb998861982a03bbe41e204", - "nvidia_cutlass_dsl-4.6.1-py3-none-any.whl": "93135a9d48e1bedf584828e0a021f174ac31591ef21f6ea53c206169ccbfab26", - "nvidia_cutlass_dsl_libs_base-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl": "dcbbf471839801501030f1097ce13dbf5082ad34f26d6b503459c6fed078e4e9", - "nvidia_cutlass_dsl_libs_core-4.6.1-py3-none-any.whl": "f1d895ee24b1ba711b2b9d4a43c62fa1f6fe3a50634e25eeffe88fa17f6c5e47", - "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl": "a8483778ad75ae50efd5b8981adfe8d7bc0c9cbcf5fd16a02b3bd9a062b98e2d" + "nvidia_cuda_nvrtc-13.2.78-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl": "a9049031da08cbedd0c20e3470e5a978dc330af0e0326b3b05774718c665dc3e", + "nvidia_cutlass_dsl-4.6.0-py3-none-any.whl": "e3e0e4d8df20d82c8401fa013f4d82021f41daa5fca3d24b55d4a677f2308ca8", + "nvidia_cutlass_dsl_libs_base-4.6.0-cp311-cp311-manylinux_2_28_x86_64.whl": "90a3e7a61d110a8ed005aae83869c6e5dca0723e36298297c0780e21db59c016", + "nvidia_cutlass_dsl_libs_core-4.6.0-py3-none-any.whl": "f9ea6d313a03cb11fa177da32e8747ad0cac51358850810f36aa6c4736192c27", + "nvidia_cutlass_dsl_libs_cu13-4.6.0-cp311-cp311-manylinux_2_28_x86_64.whl": "7235c5cfd1db3814bf559d6b976beb759dc7298914f4ff2684a91c3071369598" } } diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tests/python/tirx/iket/test_iket_orchestration.py index d72446716f4b..be12799f3acf 100644 --- a/tests/python/tirx/iket/test_iket_orchestration.py +++ b/tests/python/tirx/iket/test_iket_orchestration.py @@ -132,7 +132,7 @@ def test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey result = iket.profile( (sys.executable, "workload.py"), output_dir=target, - profile_name="cutlass-4.6.1", + profile_name="cutlass-4.6.0", postprocess="all", clobber=False, cwd=cwd, @@ -148,7 +148,7 @@ def test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey assert captured["cwd"] == cwd assert captured["timeout"] == 12.5 assert captured["env"]["IKET_TEST_ENV"] == "present" - assert captured["env"]["TVM_IKET_OFFICIAL_PROFILE"] == "cutlass-4.6.1" + assert captured["env"]["TVM_IKET_OFFICIAL_PROFILE"] == "cutlass-4.6.0" assert os.environ["TVM_IKET_OFFICIAL_PROFILE"] == "inherited-profile" assert result.output_dir == target assert result.command == (sys.executable, "workload.py") diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 956f7f34bc58..78d766045155 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -36,7 +36,7 @@ from tvm.tirx.cuda.iket import IketProfiler TARGET = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) -ORACLE_PATH = Path(__file__).parent / "oracle" / "iket_official_cutlass_4_6_1_oracle.json" +ORACLE_PATH = Path(__file__).parent / "oracle" / "iket_official_cutlass_4_6_0_oracle.json" @T.prim_func @@ -182,6 +182,23 @@ def while_carried_divergent_token(out: T.Buffer((32,), "int32")): out[tx] = guard[0] +@T.prim_func +def annotated_outer_loop_with_nested_break(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + value = T.alloc_local((1,), "int32") + value[0] = 0 + for _outer in T.serial(2, unroll=False): + iket.range_push("outer") + for inner in T.serial(2, unroll=False): + if inner == 1: + break + value[0] = value[0] + 1 + iket.range_pop() + out[tx] = value[0] + + @T.prim_func def marks_30(out: T.Buffer((32,), "int32")): T.device_entry() @@ -435,6 +452,23 @@ def test_explicit_cuda_shuffle_broadcast_is_warp_uniform(): assert "__iket_evt_decl_warp_zero_1_attrs" in source +def test_nested_loop_control_does_not_reject_annotated_outer_loop(): + source = _cuda_source(_compile(annotated_outer_loop_with_nested_break)) + assert "__iket_evt_decl_outer_1_attrs" in source + + +@pytest.mark.parametrize( + "kernel_func", (loop_carried_divergent_token, while_carried_divergent_token) +) +def test_unproven_warp_convergence_warns_and_compiles(kernel_func, capfd): + source = _cuda_source(_compile(kernel_func)) + warning = capfd.readouterr().err + + assert "IKET warp convergence could not be proven for event site" in warning + assert "continuing because convergence diagnostics are advisory" in warning + assert "__iket_evt_decl" in source + + @pytest.mark.parametrize( ("kernel_func", "message"), [ @@ -442,19 +476,9 @@ def test_explicit_cuda_shuffle_broadcast_is_warp_uniform(): pytest.param(overlapping_ranges, "strictly alternating", id="overlap"), pytest.param(repeated_range_end, "strictly alternating", id="repeated-end"), pytest.param(unbalanced_stack, "balanced range_push/range_pop", id="unbalanced-stack"), - pytest.param( - loop_carried_divergent_token, - "divergent set of lanes", - id="loop-fixed-point-divergence", - ), - pytest.param( - while_carried_divergent_token, - "divergent set of lanes", - id="while-fixed-point-divergence", - ), ], ) -def test_rejects_unsupported_or_divergent_semantics(kernel_func, message): +def test_rejects_unsupported_semantics(kernel_func, message): with pytest.raises(ValueError, match=message): _compile(kernel_func) @@ -519,7 +543,8 @@ def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch): injection_path.write_bytes(b"locked") injection_relative = "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so" profile = { - "versions": {"nvidia-cutlass-dsl-libs-base": "4.6.1"}, + "nvrtc_version": (13, 2), + "versions": {"nvidia-cutlass-dsl-libs-base": "4.6.0"}, "files": { "nvidia-cutlass-dsl-libs-base": { injection_relative: hashlib.sha256(b"locked").hexdigest() @@ -528,20 +553,20 @@ def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch): } class FakeDistribution: - version = "4.6.1" + version = "4.6.0" @staticmethod def locate_file(_relative_path): return injection_path - monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.1", profile) + monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.0", profile) monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name: FakeDistribution()) monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint", lambda: None) monkeypatch.setattr( _iket_official, "_validate_injection_environment", lambda _expected_digest: None ) - monkeypatch.setattr(_iket_official, "_validate_nvrtc_13_3", lambda: None) - monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.1") + monkeypatch.setattr(_iket_official, "_validate_nvrtc_version", lambda _version: None) + monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.0") _iket_official.validate_official_environment() monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE") @@ -584,10 +609,10 @@ def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch ) -def test_cutlass_4_6_1_oracle_manifest_integrity(): +def test_cutlass_4_6_0_oracle_manifest_integrity(): oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8")) assert oracle["schema_version"] == 2 - assert oracle["profile"]["cutlass_dsl"] == "4.6.1" + assert oracle["profile"]["cutlass_dsl"] == "4.6.0" assert oracle["profile"]["instrument_method"] == "NativeDump" assert "--dump-dir=" in oracle["profile"]["compiler_flags"] metadata_bytes = json.dumps(oracle["metadata"], sort_keys=True, separators=(",", ":")).encode() From 6dc458680f000985f507a7193bd3c0965d8d647a Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Mon, 3 Aug 2026 16:59:19 -0400 Subject: [PATCH 2/5] feat(lower-tirx): support official IKET payloads Implement payload-aware NativeDump and ExtendedNativeDump metadata and event lowering from the locked CUTLASS 4.6.0 oracle. Enable lowering only for validated injected children and keep the pass focused on ABI validation and mechanical IR rewriting without token lifetime, stack balance, or warp convergence proofs. --- python/tvm/backend/cuda/iket.py | 38 +- python/tvm/backend/cuda/intrinsics/misc.py | 5 +- python/tvm/backend/cuda/op.py | 22 +- src/backend/cuda/transforms/lower_iket.cc | 1164 +++----- .../python/tirx/iket/iket_profile_workload.py | 83 +- .../oracle/generate_iket_official_oracle.py | 326 ++- .../iket_official_cutlass_4_6_0_oracle.json | 2495 ++++++++++++++++- .../tirx/iket/test_iket_orchestration.py | 1 + tests/python/tirx/iket/test_iket_profiler.py | 528 +++- 9 files changed, 3619 insertions(+), 1043 deletions(-) diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py index c21d56835184..c1edbed2effb 100644 --- a/python/tvm/backend/cuda/iket.py +++ b/python/tvm/backend/cuda/iket.py @@ -52,6 +52,7 @@ from tvm.script import tirx as T _PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE" +_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") @@ -368,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()) @@ -455,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 diff --git a/python/tvm/backend/cuda/intrinsics/misc.py b/python/tvm/backend/cuda/intrinsics/misc.py index b6bbb9780854..816c2da3727c 100644 --- a/python/tvm/backend/cuda/intrinsics/misc.py +++ b/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", ) diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py index 7ab78c102bb0..6e56cb1fb4b6 100644 --- a/python/tvm/backend/cuda/op.py +++ b/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) diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc index ccb7e9db2435..b9ca652efa70 100644 --- a/src/backend/cuda/transforms/lower_iket.cc +++ b/src/backend/cuda/transforms/lower_iket.cc @@ -27,13 +27,13 @@ #include #include #include -#include #include #include #include #include #include +#include #include #include #include @@ -49,13 +49,16 @@ namespace transform { namespace { -constexpr int kMaxDeclarations = 30; +constexpr uint32_t kNativeMaxDeclarations = 30; +constexpr uint32_t kExtendedMaxDeclarations = 4032; +constexpr uint32_t kNativeFirstEventId = 1; +constexpr uint32_t kExtendedFirstEventId = 64; +constexpr uint32_t kRangePopEventId = 31; +constexpr uint32_t kNativeMaxEventId = 31; +constexpr uint32_t kExtendedMaxEventId = 4095; constexpr int kOfficialMetaInfoBytes = 48; constexpr int kOfficialEventAttributesBytes = 60; constexpr int kOfficialRangeAttributesBytes = 72; -constexpr size_t kMaxTokenAnalysisStates = 256; -constexpr size_t kMaxTokenAnalysisIterations = 256; -constexpr size_t kMaxConvergenceAnalysisIterations = 256; enum class DeclarationKind : int { kRange = 1, @@ -63,6 +66,25 @@ enum class DeclarationKind : int { kMark = 3, }; +enum class InstrumentMode : uint32_t { + kNativeDump = 3, + kExtendedNativeDump = 5, +}; + +enum class PayloadType : uint32_t { + kNone = 0, + kI8 = 1, + kUI8 = 2, + kI16 = 3, + kUI16 = 4, + kI32 = 5, + kUI32 = 6, + kI64 = 7, + kFP32 = 13, + kFP64 = 14, + kUI64 = 16, +}; + const Op& IketMarkOp() { static const Op& op = Op::Get("tirx.cuda.iket_mark"); return op; @@ -157,19 +179,80 @@ std::string GetName(const CallNode* call, size_t index = 0) { return result; } -std::string ValidatePayload(const Expr& payload) { - DLDataType dtype = payload.as_or_throw().ty()->dtype; +PayloadType ValidatePayload(const Expr& payload) { + ffi::Optional value_optional = payload.as(); + TVM_FFI_CHECK(value_optional.has_value(), TypeError) + << "IKET payload must be a scalar numeric value, but got " << payload; + PrimExpr value = value_optional.value(); + ffi::Optional type = value->ty.as(); + TVM_FFI_CHECK(type.has_value(), TypeError) + << "IKET payload must be a scalar numeric value, but got " << value->ty; + DLDataType dtype = type.value()->dtype; TVM_FFI_CHECK_EQ(dtype.lanes, 1, TypeError) << "IKET payload must be a scalar value"; - TVM_FFI_CHECK_GT(dtype.bits, 0, TypeError) << "IKET payload must have a concrete bit width"; - TVM_FFI_CHECK_LE(dtype.bits, 64, TypeError) << "IKET payload must be at most 64 bits"; - bool supported = dtype.code == kDLInt || dtype.code == kDLUInt || dtype.code == kDLFloat || - dtype.code == kDLBfloat; - TVM_FFI_CHECK(supported, TypeError) << "IKET payload must be bool, int, uint, or float, but got " - << ffi::DLDataTypeToString(dtype); - CallEffectKind effect = SideEffect(payload.as_or_throw()); + PayloadType payload_type; + if (dtype.code == kDLBool) { + payload_type = PayloadType::kUI8; + } else if (dtype.code == kDLInt && dtype.bits == 8) { + payload_type = PayloadType::kI8; + } else if (dtype.code == kDLUInt && dtype.bits == 8) { + payload_type = PayloadType::kUI8; + } else if (dtype.code == kDLInt && dtype.bits == 16) { + payload_type = PayloadType::kI16; + } else if (dtype.code == kDLUInt && dtype.bits == 16) { + payload_type = PayloadType::kUI16; + } else if (dtype.code == kDLInt && dtype.bits == 32) { + payload_type = PayloadType::kI32; + } else if (dtype.code == kDLUInt && dtype.bits == 32) { + payload_type = PayloadType::kUI32; + } else if (dtype.code == kDLInt && dtype.bits == 64) { + payload_type = PayloadType::kI64; + } else if (dtype.code == kDLUInt && dtype.bits == 64) { + payload_type = PayloadType::kUI64; + } else if (dtype.code == kDLFloat && dtype.bits == 32) { + payload_type = PayloadType::kFP32; + } else if (dtype.code == kDLFloat && dtype.bits == 64) { + payload_type = PayloadType::kFP64; + } else { + TVM_FFI_THROW(TypeError) + << "IKET payload supports only bool, int/uint 8/16/32/64, float32, and float64; got " + << ffi::DLDataTypeToString(dtype); + } + CallEffectKind effect = SideEffect(value); TVM_FFI_CHECK(effect <= CallEffectKind::kReadState, ValueError) << "IKET payload expressions may read state but must not update state; got " << effect; - return ffi::DLDataTypeToString(dtype); + return payload_type; +} + +const char* PayloadTypeName(PayloadType type) { + switch (type) { + case PayloadType::kNone: + return "NoPayload"; + case PayloadType::kI8: + return "I8"; + case PayloadType::kUI8: + return "UI8"; + case PayloadType::kI16: + return "I16"; + case PayloadType::kUI16: + return "UI16"; + case PayloadType::kI32: + return "I32"; + case PayloadType::kUI32: + return "UI32"; + case PayloadType::kI64: + return "I64"; + case PayloadType::kFP32: + return "FP32"; + case PayloadType::kFP64: + return "FP64"; + case PayloadType::kUI64: + return "UI64"; + } + return "Unknown"; +} + +bool Is64BitPayload(PayloadType type) { + return type == PayloadType::kI64 || type == PayloadType::kUI64 || type == PayloadType::kFP64; } bool IsScalarBufferAccess(const BufferVar& buffer, const ffi::Array& indices) { @@ -204,10 +287,10 @@ struct Declaration { DeclarationKey key; bool payload_schema_set{false}; bool has_payload{false}; - std::string payload_dtype; + PayloadType payload_type{PayloadType::kNone}; bool end_payload_schema_set{false}; bool end_has_payload{false}; - std::string end_payload_dtype; + PayloadType end_payload_type{PayloadType::kNone}; uint32_t event_id{0}; }; @@ -215,14 +298,16 @@ class AnnotationCollector : public StmtExprVisitor { public: std::map declarations; bool has_annotations{false}; + bool has_payload_calls{false}; private: - void AddDeclaration(DeclarationKind kind, const CallNode* call, bool sentinel = false) { + void AddDeclaration(DeclarationKind kind, const CallNode* call) { std::string name = GetName(call); TVM_FFI_CHECK(call->args.size() == 1 || call->args.size() == 2, TypeError) << call->op.as().value()->name << " expects a name and optional payload"; bool has_payload = call->args.size() == 2; - std::string dtype = has_payload ? ValidatePayload(call->args[1]) : std::string(); + has_payload_calls = has_payload_calls || has_payload; + PayloadType payload_type = has_payload ? ValidatePayload(call->args[1]) : PayloadType::kNone; auto [name_it, name_inserted] = declaration_kinds_.emplace(name, kind); TVM_FFI_CHECK(name_inserted || name_it->second == kind, ValueError) << "IKET declaration " << name << " changes event kind from " @@ -230,19 +315,18 @@ class AnnotationCollector : public StmtExprVisitor { DeclarationKey key{kind, std::move(name)}; auto [it, inserted] = declarations.emplace(key, Declaration{key}); Declaration& declaration = it->second; - if (!sentinel) { - if (declaration.payload_schema_set) { - TVM_FFI_CHECK_EQ(declaration.has_payload, has_payload, ValueError) - << "IKET declaration " << declaration.key.name - << " changes payload presence between sites"; - TVM_FFI_CHECK_EQ(declaration.payload_dtype, dtype, TypeError) - << "IKET declaration " << declaration.key.name - << " changes payload dtype between sites"; - } else { - declaration.payload_schema_set = true; - declaration.has_payload = has_payload; - declaration.payload_dtype = std::move(dtype); - } + if (declaration.payload_schema_set) { + TVM_FFI_CHECK_EQ(declaration.has_payload, has_payload, ValueError) + << "IKET declaration " << declaration.key.name + << " changes payload presence between sites"; + TVM_FFI_CHECK(declaration.payload_type == payload_type, TypeError) + << "IKET declaration " << declaration.key.name << " changes payload type from " + << PayloadTypeName(declaration.payload_type) << " to " << PayloadTypeName(payload_type) + << " between sites"; + } else { + declaration.payload_schema_set = true; + declaration.has_payload = has_payload; + declaration.payload_type = payload_type; } } @@ -267,7 +351,9 @@ class AnnotationCollector : public StmtExprVisitor { call->ty.as_or_throw()->dtype.bits == 32, TypeError) << "IKET sentinel_token must return uint32"; - AddDeclaration(DeclarationKind::kRange, call, true); + // A sentinel carries only token-flow identity. It emits no runtime + // event and therefore must not create metadata or consume an event ID. + GetName(call); } else if (call->op.same_as(IketRangePushOp())) { AddDeclaration(DeclarationKind::kPush, call); } else if (call->op.same_as(IketRangeEndOp())) { @@ -277,7 +363,10 @@ class AnnotationCollector : public StmtExprVisitor { TVM_FFI_CHECK(token_dtype.code == kDLUInt && token_dtype.bits == 32 && token_dtype.lanes == 1, TypeError) << "IKET RangeToken must have dtype uint32"; - if (call->args.size() == 2) ValidatePayload(call->args[1]); + if (call->args.size() == 2) { + has_payload_calls = true; + ValidatePayload(call->args[1]); + } } else if (call->op.same_as(IketRangePopOp())) { TVM_FFI_CHECK_EQ(call->args.size(), 0, TypeError) << "IKET range_pop takes no arguments"; } @@ -378,20 +467,25 @@ class RangeEndSchemaVerifier : public StmtExprVisitor { auto possible_it = token_declarations_.find(token->buffer.get()); TVM_FFI_ICHECK(possible_it != token_declarations_.end()); bool has_payload = call->args.size() == 2; - std::string dtype = has_payload ? ValidatePayload(call->args[1]) : std::string(); + PayloadType payload_type = has_payload ? ValidatePayload(call->args[1]) : PayloadType::kNone; for (const DeclarationKey& key : possible_it->second) { auto declaration_it = declarations_->find(key); - TVM_FFI_ICHECK(declaration_it != declarations_->end()); + // A sentinel-only identity has no real declaration and no payload + // schema. If the same token can also carry a real start, that real + // declaration is still checked below. + if (declaration_it == declarations_->end()) continue; Declaration& declaration = declaration_it->second; if (declaration.end_payload_schema_set) { TVM_FFI_CHECK_EQ(declaration.end_has_payload, has_payload, ValueError) << "range_end for " << key.name << " changes payload presence between sites"; - TVM_FFI_CHECK_EQ(declaration.end_payload_dtype, dtype, TypeError) - << "range_end for " << key.name << " changes payload dtype between sites"; + TVM_FFI_CHECK(declaration.end_payload_type == payload_type, TypeError) + << "range_end for " << key.name << " changes payload type from " + << PayloadTypeName(declaration.end_payload_type) << " to " + << PayloadTypeName(payload_type) << " between sites"; } else { declaration.end_payload_schema_set = true; declaration.end_has_payload = has_payload; - declaration.end_payload_dtype = dtype; + declaration.end_payload_type = payload_type; } } StmtExprVisitor::VisitExpr_(call); @@ -401,6 +495,19 @@ class RangeEndSchemaVerifier : public StmtExprVisitor { std::map* declarations_; }; +void ValidateRangeSchemas(const std::map& declarations) { + for (const auto& [key, declaration] : declarations) { + if (key.kind != DeclarationKind::kRange || !declaration.end_payload_schema_set) continue; + TVM_FFI_CHECK_EQ(declaration.has_payload, declaration.end_has_payload, ValueError) + << "IKET token range " << key.name + << " must use payloads at both range_start and range_end, or at neither endpoint"; + TVM_FFI_CHECK(declaration.payload_type == declaration.end_payload_type, TypeError) + << "IKET token range " << key.name << " changes payload type from " + << PayloadTypeName(declaration.payload_type) << " at range_start to " + << PayloadTypeName(declaration.end_payload_type) << " at range_end"; + } +} + class TokenVerifier : public StmtExprVisitor { public: explicit TokenVerifier(const TokenBufferSet& token_buffers) : token_buffers_(token_buffers) {} @@ -474,370 +581,6 @@ class TokenVerifier : public StmtExprVisitor { bool allow_producer_{false}; }; -enum class TokenValueKind : uint8_t { - kSentinel, - kActiveRange, - kConsumed, -}; - -struct TokenValue { - TokenValueKind kind{TokenValueKind::kSentinel}; - std::string name; - - bool operator==(const TokenValue& other) const { - return kind == other.kind && name == other.name; - } - - bool operator<(const TokenValue& other) const { - return std::tie(kind, name) < std::tie(other.kind, other.name); - } -}; - -struct TokenAnalysisState { - std::map token_values; - std::set active_ranges; - - bool operator==(const TokenAnalysisState& other) const { - return token_values == other.token_values && active_ranges == other.active_ranges; - } - - bool operator<(const TokenAnalysisState& other) const { - return std::tie(token_values, active_ranges) < - std::tie(other.token_values, other.active_ranges); - } -}; - -using TokenAnalysisStates = std::set; - -class TokenOperationFinder : public StmtExprVisitor { - public: - bool found{false}; - - private: - void VisitExpr_(const CallNode* call) final { - if (IsTokenProducer(call) || call->op.same_as(IketRangeEndOp())) { - found = true; - return; - } - StmtExprVisitor::VisitExpr_(call); - } -}; - -/*! \brief Prove the strict token alternation required by NVIDIA IKET. - * - * Any state explosion, unsupported control flow, unknown token, second start of - * an active name, or end of an inactive name rejects the annotated module. - */ -class OfficialTokenAnalyzer { - public: - explicit OfficialTokenAnalyzer(const TokenBufferSet& token_buffers) - : token_buffers_(token_buffers) {} - - bool Prove(const Stmt& body) { - TokenAnalysisStates initial{TokenAnalysisState{}}; - TokenAnalysisStates output = Process(body, initial); - if (!valid_) return false; - for (const TokenAnalysisState& state : output) { - if (!state.active_ranges.empty()) return false; - } - return true; - } - - private: - TokenAnalysisStates Limit(TokenAnalysisStates states) { - if (states.size() > kMaxTokenAnalysisStates) valid_ = false; - return valid_ ? std::move(states) : TokenAnalysisStates{}; - } - - TokenAnalysisStates Union(const TokenAnalysisStates& lhs, const TokenAnalysisStates& rhs) { - TokenAnalysisStates result = lhs; - result.insert(rhs.begin(), rhs.end()); - return Limit(std::move(result)); - } - - TokenAnalysisStates ProcessLoop(const Stmt& body, const TokenAnalysisStates& input) { - TokenAnalysisStates closure = input; - for (size_t iteration = 0; valid_ && iteration < kMaxTokenAnalysisIterations; ++iteration) { - TokenAnalysisStates after_body = Process(body, closure); - if (!valid_) return {}; - TokenAnalysisStates next = Union(closure, after_body); - if (!valid_) return {}; - if (next == closure) return closure; - closure = std::move(next); - } - valid_ = false; - return {}; - } - - TokenAnalysisStates ProcessTokenStore(const BufferStoreNode* store, - const TokenAnalysisStates& input) { - TokenAnalysisStates result; - for (const TokenAnalysisState& old_state : input) { - TokenAnalysisState state = old_state; - if (const auto* call = store->value.as(); call && IsTokenProducer(call)) { - if (call->op.same_as(IketSentinelOp())) { - state.token_values[store->buffer.get()] = TokenValue{TokenValueKind::kSentinel, {}}; - } else { - std::string name = GetName(call); - if (state.active_ranges.count(name)) { - valid_ = false; - return {}; - } - state.active_ranges.insert(name); - state.token_values[store->buffer.get()] = - TokenValue{TokenValueKind::kActiveRange, std::move(name)}; - } - } else if (const auto* load = store->value.as(); - load && token_buffers_.count(load->buffer.get())) { - auto source = state.token_values.find(load->buffer.get()); - if (source == state.token_values.end()) { - valid_ = false; - return {}; - } - state.token_values[store->buffer.get()] = source->second; - } else { - valid_ = false; - return {}; - } - result.insert(std::move(state)); - if (store->predicate.has_value()) result.insert(old_state); - } - return Limit(std::move(result)); - } - - TokenAnalysisStates ProcessRangeEnd(const CallNode* call, const TokenAnalysisStates& input) { - const auto* load = call->args[0].as(); - if (!load || !token_buffers_.count(load->buffer.get())) { - valid_ = false; - return {}; - } - TokenAnalysisStates result; - for (const TokenAnalysisState& old_state : input) { - auto token = old_state.token_values.find(load->buffer.get()); - if (token == old_state.token_values.end()) { - valid_ = false; - return {}; - } - if (token->second.kind == TokenValueKind::kConsumed) { - valid_ = false; - return {}; - } - TokenAnalysisState state = old_state; - if (token->second.kind == TokenValueKind::kActiveRange) { - auto active = state.active_ranges.find(token->second.name); - if (active == state.active_ranges.end()) { - valid_ = false; - return {}; - } - state.active_ranges.erase(active); - } - state.token_values[load->buffer.get()] = TokenValue{TokenValueKind::kConsumed, {}}; - result.insert(std::move(state)); - } - return Limit(std::move(result)); - } - - TokenAnalysisStates Process(const Stmt& stmt, const TokenAnalysisStates& input) { - if (!valid_ || input.empty()) return input; - if (const auto* sequence = stmt.as()) { - TokenAnalysisStates states = input; - for (const Stmt& item : sequence->seq) { - states = Process(item, states); - if (!valid_ || states.empty()) break; - } - return states; - } - if (const auto* store = stmt.as()) { - if (token_buffers_.count(store->buffer.get())) { - return ProcessTokenStore(store, input); - } - return input; - } - if (const auto* evaluate = stmt.as()) { - if (const auto* call = evaluate->value.as()) { - if (call->op.same_as(IketRangeEndOp())) return ProcessRangeEnd(call, input); - if (call->op.same_as(builtin::thread_return())) { - for (const TokenAnalysisState& state : input) { - if (!state.active_ranges.empty()) { - valid_ = false; - break; - } - } - return {}; - } - } - return input; - } - if (const auto* branch = stmt.as()) { - TokenAnalysisStates then_states = Process(branch->then_case, input); - TokenAnalysisStates else_states = - branch->else_case.has_value() ? Process(branch->else_case.value(), input) : input; - return Union(then_states, else_states); - } - if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); - if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); - if (const auto* attr_stmt = stmt.as()) { - return Process(attr_stmt->body, input); - } - if (const auto* block = stmt.as()) { - TokenAnalysisStates states = input; - if (block->init.has_value()) states = Union(states, Process(block->init.value(), states)); - return Process(block->body, states); - } - if (const auto* realize = stmt.as()) { - return Union(input, Process(realize->block, input)); - } - if (stmt.as() || stmt.as()) { - valid_ = false; - return {}; - } - // Unknown statement forms are harmless only if they cannot hide token - // operations. Otherwise their control flow has not been proved. - TokenOperationFinder finder; - finder(stmt); - if (finder.found) { - valid_ = false; - return {}; - } - return input; - } - - const TokenBufferSet& token_buffers_; - bool valid_{true}; -}; - -using OfficialStackState = std::vector; -using OfficialStackStates = std::set; - -class StackOperationFinder : public StmtExprVisitor { - public: - bool found{false}; - - private: - void VisitExpr_(const CallNode* call) final { - if (call->op.same_as(IketRangePushOp()) || call->op.same_as(IketRangePopOp())) { - found = true; - return; - } - StmtExprVisitor::VisitExpr_(call); - } -}; - -/*! \brief Prove balanced LIFO push/pop behavior required by NVIDIA IKET. */ -class OfficialStackAnalyzer { - public: - bool Prove(const Stmt& body) { - OfficialStackStates output = Process(body, OfficialStackStates{OfficialStackState{}}); - if (!valid_) return false; - for (const OfficialStackState& state : output) { - if (!state.empty()) return false; - } - return true; - } - - private: - OfficialStackStates Limit(OfficialStackStates states) { - if (states.size() > kMaxTokenAnalysisStates) valid_ = false; - return valid_ ? std::move(states) : OfficialStackStates{}; - } - - OfficialStackStates Union(const OfficialStackStates& lhs, const OfficialStackStates& rhs) { - OfficialStackStates result = lhs; - result.insert(rhs.begin(), rhs.end()); - return Limit(std::move(result)); - } - - OfficialStackStates ProcessLoop(const Stmt& body, const OfficialStackStates& input) { - OfficialStackStates closure = input; - for (size_t iteration = 0; valid_ && iteration < kMaxTokenAnalysisIterations; ++iteration) { - OfficialStackStates after_body = Process(body, closure); - if (!valid_) return {}; - OfficialStackStates next = Union(closure, after_body); - if (!valid_) return {}; - if (next == closure) return closure; - closure = std::move(next); - } - valid_ = false; - return {}; - } - - OfficialStackStates Process(const Stmt& stmt, const OfficialStackStates& input) { - if (!valid_ || input.empty()) return input; - if (const auto* sequence = stmt.as()) { - OfficialStackStates states = input; - for (const Stmt& item : sequence->seq) { - states = Process(item, states); - if (!valid_ || states.empty()) break; - } - return states; - } - if (const auto* evaluate = stmt.as()) { - if (const auto* call = evaluate->value.as()) { - if (call->op.same_as(IketRangePushOp())) { - OfficialStackStates result; - for (OfficialStackState state : input) { - state.push_back(GetName(call)); - result.insert(std::move(state)); - } - return Limit(std::move(result)); - } - if (call->op.same_as(IketRangePopOp())) { - OfficialStackStates result; - for (OfficialStackState state : input) { - if (state.empty()) { - valid_ = false; - return {}; - } - state.pop_back(); - result.insert(std::move(state)); - } - return Limit(std::move(result)); - } - if (call->op.same_as(builtin::thread_return())) { - for (const OfficialStackState& state : input) { - if (!state.empty()) { - valid_ = false; - break; - } - } - return {}; - } - } - return input; - } - if (const auto* branch = stmt.as()) { - OfficialStackStates then_states = Process(branch->then_case, input); - OfficialStackStates else_states = - branch->else_case.has_value() ? Process(branch->else_case.value(), input) : input; - return Union(then_states, else_states); - } - if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); - if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); - if (const auto* attr_stmt = stmt.as()) return Process(attr_stmt->body, input); - if (const auto* block = stmt.as()) { - OfficialStackStates states = input; - if (block->init.has_value()) states = Union(states, Process(block->init.value(), states)); - return Process(block->body, states); - } - if (const auto* realize = stmt.as()) { - return Union(input, Process(realize->block, input)); - } - if (stmt.as() || stmt.as()) { - valid_ = false; - return {}; - } - StackOperationFinder finder; - finder(stmt); - if (finder.found) { - valid_ = false; - return {}; - } - return input; - } - - bool valid_{true}; -}; - class StripIket : public StmtExprMutator { public: explicit StripIket(TokenBufferSet token_buffers) : token_buffers_(std::move(token_buffers)) {} @@ -953,17 +696,6 @@ bool IsSm90OrNewer(const PrimFunc& function) { return std::stoi(value.substr(3, end - 3)) >= 90; } -bool HasAnyPayload(const std::map& declarations) { - for (const auto& item : declarations) { - const Declaration& declaration = item.second; - if (declaration.has_payload || - (declaration.end_payload_schema_set && declaration.end_has_payload)) { - return true; - } - } - return false; -} - std::string FunctionName(const GlobalVar& global_var, const PrimFunc& function) { if (auto symbol = function->GetAttr(tvm::attr::kGlobalSymbol)) { return symbol.value(); @@ -976,6 +708,7 @@ struct KernelIketInfo { PrimFunc function; std::string name; std::map declarations; + bool has_payload_calls{false}; }; void PutOfficialU32(std::vector* bytes, size_t offset, uint32_t value) { @@ -1028,8 +761,11 @@ std::map CollectOfficialDeclarations( return declarations; } -std::string BuildOfficialDeviceSource(const std::vector& kernels) { +std::string BuildOfficialDeviceSource(const std::vector& kernels, + InstrumentMode mode) { std::map declarations = CollectOfficialDeclarations(kernels); + uint32_t max_event_id = + mode == InstrumentMode::kNativeDump ? kNativeMaxEventId : kExtendedMaxEventId; std::ostringstream os; os << R"IKET( extern "C" { @@ -1039,7 +775,7 @@ extern "C" { PutOfficialU32(&meta, 0, kOfficialMetaInfoBytes); PutOfficialU32(&meta, 4, 0); PutOfficialU32(&meta, 8, 5); - PutOfficialU32(&meta, 12, 31); + PutOfficialU32(&meta, 12, max_event_id); PutOfficialU32(&meta, 16, 32); PutOfficialU32(&meta, 20, kOfficialEventAttributesBytes); PutOfficialU32(&meta, 24, 0xbabef19dU); @@ -1059,8 +795,8 @@ extern "C" { std::vector event(kOfficialEventAttributesBytes); PutOfficialU32(&event, 0, kOfficialEventAttributesBytes); PutOfficialU32(&event, 4, declaration.event_id); - PutOfficialU32(&event, 8, 3); - PutOfficialU32(&event, 12, 0); + PutOfficialU32(&event, 8, static_cast(mode)); + PutOfficialU32(&event, 12, static_cast(declaration.payload_type)); uint32_t event_position = key.kind == DeclarationKind::kRange ? 4 : (key.kind == DeclarationKind::kPush ? 1 : 0); PutOfficialU32(&event, 16, event_position); @@ -1086,7 +822,14 @@ extern "C" { os << OfficialByteArray(range_symbol, range); } } - os << R"IKET( + bool has_payload_calls = + std::any_of(kernels.begin(), kernels.end(), + [](const KernelIketInfo& kernel) { return kernel.has_payload_calls; }); + if (mode == InstrumentMode::kNativeDump && !has_payload_calls) { + // Keep this source byte-for-byte compatible with the original NativeDump + // no-payload helper. Existing MegaKernel placeholders and SASS must not + // move merely because payload and ExtendedNativeDump support exists. + os << R"IKET( } template @@ -1110,12 +853,12 @@ __forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( unsigned int event_id) { switch (event_id) { )IKET"; - for (const auto& [key, declaration] : declarations) { - os << " case " << declaration.event_id << ":\n" - << " tvm_builtin_iket_official_event_impl<" << declaration.event_id << ">();\n" - << " break;\n"; - } - os << R"IKET( case 31: + for (const auto& [key, declaration] : declarations) { + os << " case " << declaration.event_id << ":\n" + << " tvm_builtin_iket_official_event_impl<" << declaration.event_id << ">();\n" + << " break;\n"; + } + os << R"IKET( case 31: tvm_builtin_iket_official_event_impl<31>(); break; default: @@ -1124,285 +867,234 @@ __forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( return event_id; } )IKET"; - return os.str(); -} - -using UniformBufferSet = std::unordered_set; -using DivergentVarSet = std::unordered_set; - -UniformBufferSet IntersectUniformBuffers(const UniformBufferSet& lhs, const UniformBufferSet& rhs) { - UniformBufferSet result; - for (const VarNode* buffer : lhs) { - if (rhs.count(buffer)) result.insert(buffer); + return os.str(); } - return result; -} -DivergentVarSet UnionDivergentVars(const DivergentVarSet& lhs, const DivergentVarSet& rhs) { - DivergentVarSet result = lhs; - result.insert(rhs.begin(), rhs.end()); - return result; + os << R"IKET( } +)IKET"; -class UniformExprChecker : public ExprVisitor { - public: - UniformExprChecker(const DivergentVarSet& divergent_vars, const UniformBufferSet& uniform_buffers) - : divergent_vars_(divergent_vars), uniform_buffers_(uniform_buffers) {} - - bool IsUniform(const Expr& expr) { - uniform_ = true; - operator()(expr); - return uniform_; - } - - private: - void VisitExpr_(const VarNode* var) final { - if (divergent_vars_.count(var)) uniform_ = false; - } - - void VisitExpr_(const BufferLoadNode* load) final { - if (!uniform_buffers_.count(load->buffer.get()) || - !IsScalarBufferAccess(load->buffer, load->indices)) { - uniform_ = false; - return; - } - for (const PrimExpr& index : load->indices) VisitExpr(index); - } - - void VisitExpr_(const CallNode* call) final { - // Calls may read lane-local or device state even when their explicit - // arguments are uniform. Treat them conservatively, except for likely(), - // which is only an annotation around its argument. - if (call->op.same_as(builtin::likely())) { - ExprVisitor::VisitExpr_(call); - } else if (IsTokenProducer(call)) { - // Both producers return a declaration id (or sentinel zero) that is - // independent of a potentially lane-varying payload. - return; - } else if (call->op.same_as(builtin::tvm_warp_shuffle()) && call->args.size() == 5) { - // A shuffle from one warp-uniform source lane is a broadcast. Its - // value may depend on threadIdx, but every active lane observes the - // selected lane's value. - VisitExpr(call->args[0]); - for (size_t i = 2; i < call->args.size(); ++i) VisitExpr(call->args[i]); - } else if (call->op.same_as(Op::Get("tirx.cuda.__shfl_sync")) && call->args.size() == 4) { - // CUDA's explicit __shfl_sync(mask, value, src_lane, width) has the - // same broadcast semantics when mask/src_lane/width are uniform. - // Ignore the lane-local value, but prove the control operands uniform. - VisitExpr(call->args[0]); - VisitExpr(call->args[2]); - VisitExpr(call->args[3]); - } else if (call->op.same_as(builtin::bitwise_and()) || - call->op.same_as(builtin::bitwise_or()) || - call->op.same_as(builtin::bitwise_xor()) || - call->op.same_as(builtin::bitwise_not())) { - // These integer/boolean operators are pure. A composed guard remains - // uniform exactly when each operand is uniform. - ExprVisitor::VisitExpr_(call); - } else { - uniform_ = false; - } - } - - const DivergentVarSet& divergent_vars_; - const UniformBufferSet& uniform_buffers_; - bool uniform_{true}; -}; + if (mode == InstrumentMode::kNativeDump) { + os << R"IKET( +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_impl() { + asm volatile( + "{\n" + ".reg .b32 %%r, %%t;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "or.b32 %%t, %%t, %0;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "st.weak.shared.u32 [%%r], %%t;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId) + : "memory"); +} -class LoopControlFinder : public StmtExprVisitor { - public: - bool found{false}; +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_payload32_impl( + unsigned int payload) { + asm volatile( + "{\n" + ".reg .pred %%p;\n" + ".reg .b32 %%r, %%t, %%mask, %%payload32;\n" + "activemask.b32 %%mask;\n" + "elect.sync _|%%p, %%mask;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "or.b32 %%t, %%t, %0;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "mov.b32 %%payload32, %1;\n" + "@%%p st.weak.shared.u32 [%%r], %%t;\n" + "@%%p st.weak.shared.b32 [%%r+4], %%payload32;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId), "r"(payload) + : "memory"); +} - private: - // A break or continue nested inside another loop targets that inner loop, - // not the loop whose body this finder was asked to inspect. Each nested - // loop is checked independently by IketConvergenceVerifier. - void VisitStmt_(const ForNode*) final {} - void VisitStmt_(const WhileNode*) final {} - void VisitStmt_(const BreakNode* op) final { found = true; } - void VisitStmt_(const ContinueNode* op) final { found = true; } - void VisitExpr_(const CallNode* call) final { - if (call->op.same_as(builtin::break_loop()) || call->op.same_as(builtin::continue_loop())) { - found = true; - return; - } - StmtExprVisitor::VisitExpr_(call); - } -}; +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_payload64_impl( + unsigned long long payload) { + asm volatile( + "{\n" + ".reg .pred %%p;\n" + ".reg .b32 %%r, %%t, %%mask;\n" + ".reg .b64 %%payload64;\n" + "activemask.b32 %%mask;\n" + "elect.sync _|%%p, %%mask;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "or.b32 %%t, %%t, %0;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "mov.b64 %%payload64, %1;\n" + "@%%p st.weak.shared.u32 [%%r], %%t;\n" + "@%%p st.weak.shared.b64 [%%r+8], %%payload64;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId), "l"(payload) + : "memory"); +} +)IKET"; + } else { + os << R"IKET( +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_impl() { + asm volatile( + "{\n" + ".reg .b32 %%r, %%t, %%evtid;\n" + ".reg .b64 %%ts_evtid;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "mov.b32 %%evtid, %0;\n" + "mov.b64 %%ts_evtid, {%%t, %%evtid};\n" + "st.weak.shared.u64 [%%r], %%ts_evtid;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId) + : "memory"); +} -class AnnotationFinder : public StmtExprVisitor { - public: - bool found{false}; +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_payload32_impl( + unsigned int payload) { + asm volatile( + "{\n" + ".reg .pred %%p;\n" + ".reg .b32 %%r, %%t, %%mask, %%evtid, %%payload32;\n" + ".reg .b64 %%ts_evtid;\n" + "activemask.b32 %%mask;\n" + "elect.sync _|%%p, %%mask;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "mov.b32 %%evtid, %0;\n" + "mov.b32 %%payload32, %1;\n" + "mov.b64 %%ts_evtid, {%%t, %%evtid};\n" + "@%%p st.weak.shared.u64 [%%r], %%ts_evtid;\n" + "@%%p st.weak.shared.b32 [%%r+8], %%payload32;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId), "r"(payload) + : "memory"); +} - private: - void VisitExpr_(const CallNode* call) final { - if (IsIketOp(call->op)) { - found = true; - return; - } - StmtExprVisitor::VisitExpr_(call); +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_payload64_impl( + unsigned long long payload) { + asm volatile( + "{\n" + ".reg .pred %%p;\n" + ".reg .b32 %%r, %%t, %%mask, %%evtid;\n" + ".reg .b64 %%ts_evtid, %%payload64;\n" + "activemask.b32 %%mask;\n" + "elect.sync _|%%p, %%mask;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "mov.b32 %%evtid, %0;\n" + "mov.b64 %%payload64, %1;\n" + "mov.b64 %%ts_evtid, {%%t, %%evtid};\n" + "@%%p st.weak.shared.u64 [%%r], %%ts_evtid;\n" + "@%%p st.weak.shared.b64 [%%r+8], %%payload64;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId), "l"(payload) + : "memory"); +} +)IKET"; } -}; -class IketConvergenceVerifier : public StmtExprVisitor { - private: - bool IsUniform(const Expr& expr) const { - return UniformExprChecker(divergent_vars_, uniform_buffers_).IsUniform(expr); + os << R"IKET( +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id) { + switch (event_id) { +)IKET"; + for (const auto& [key, declaration] : declarations) { + if (declaration.has_payload) continue; + os << " case " << declaration.event_id << ":\n" + << " tvm_builtin_iket_official_event_impl<" << declaration.event_id << ">();\n" + << " break;\n"; } - - void VisitStmt_(const BufferStoreNode* op) final { - bool uniform_store = !divergent_context_ && IsScalarBufferAccess(op->buffer, op->indices) && - IsUniform(op->value); - for (const PrimExpr& index : op->indices) { - uniform_store = uniform_store && IsUniform(index); - VisitExpr(index); - } - VisitExpr(op->value); - if (uniform_store) { - uniform_buffers_.insert(op->buffer.get()); - } else { - uniform_buffers_.erase(op->buffer.get()); - } + os << R"IKET( case 31: + tvm_builtin_iket_official_event_impl<31>(); + break; + default: + break; } + return event_id; +} - void VisitStmt_(const AttrStmtNode* op) final { - const VarNode* thread_var = nullptr; - if (op->attr_key == attr::thread_extent) { - std::string thread_tag; - if (auto iter_var = op->node.as()) { - thread_tag = iter_var.value()->thread_tag; - thread_var = iter_var.value()->var.get(); - } else if (auto var = op->node.as()) { - thread_tag = var.value()->name; - thread_var = var.value().get(); - } - if (!thread_tag.starts_with("threadIdx.")) thread_var = nullptr; - } - bool inserted = thread_var && divergent_vars_.insert(thread_var).second; - VisitExpr(op->value); - VisitStmt(op->body); - if (inserted) divergent_vars_.erase(thread_var); +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id, unsigned int payload) { + switch (event_id) { +)IKET"; + for (const auto& [key, declaration] : declarations) { + if (!declaration.has_payload || Is64BitPayload(declaration.payload_type)) continue; + os << " case " << declaration.event_id << ":\n" + << " tvm_builtin_iket_official_event_payload32_impl<" << declaration.event_id + << ">(payload);\n" + << " break;\n"; } - - void VisitStmt_(const BindNode* op) final { - if (!IsUniform(op->value)) divergent_vars_.insert(op->var.get()); - VisitExpr(op->value); + os << R"IKET( default: + break; } + return event_id; +} - void VisitStmt_(const IfThenElseNode* op) final { - VisitExpr(op->condition); - bool old_divergent = divergent_context_; - bool condition_uniform = IsUniform(op->condition); - divergent_context_ = divergent_context_ || !condition_uniform; - auto old_vars = divergent_vars_; - auto old_buffers = uniform_buffers_; - VisitStmt(op->then_case); - auto then_buffers = uniform_buffers_; - divergent_vars_ = old_vars; - uniform_buffers_ = old_buffers; - if (op->else_case.has_value()) { - VisitStmt(op->else_case.value()); - } - uniform_buffers_ = IntersectUniformBuffers(then_buffers, uniform_buffers_); - divergent_vars_ = std::move(old_vars); - divergent_context_ = old_divergent; +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id, unsigned long long payload) { + switch (event_id) { +)IKET"; + for (const auto& [key, declaration] : declarations) { + if (!declaration.has_payload || !Is64BitPayload(declaration.payload_type)) continue; + os << " case " << declaration.event_id << ":\n" + << " tvm_builtin_iket_official_event_payload64_impl<" << declaration.event_id + << ">(payload);\n" + << " break;\n"; } - - void VisitStmt_(const ForNode* op) final { - VisitExpr(op->min); - VisitExpr(op->extent); - bool uniform_loop = IsUniform(op->min) && IsUniform(op->extent); - AnnotationFinder annotations; - annotations(op->body); - LoopControlFinder loop_control; - loop_control(op->body); - TVM_FFI_CHECK(!(annotations.found && loop_control.found), ValueError) - << "IKET event sites are not allowed in a loop containing break or continue"; - - bool old_divergent = divergent_context_; - DivergentVarSet old_vars = divergent_vars_; - UniformBufferSet old_buffers = uniform_buffers_; - DivergentVarSet loop_entry_vars = old_vars; - if (!uniform_loop) loop_entry_vars.insert(op->loop_var.get()); - - DivergentVarSet loop_head_vars = loop_entry_vars; - UniformBufferSet loop_head_buffers = old_buffers; - bool converged = false; - for (size_t iteration = 0; iteration < kMaxConvergenceAnalysisIterations; ++iteration) { - divergent_vars_ = loop_head_vars; - uniform_buffers_ = loop_head_buffers; - divergent_context_ = old_divergent || !uniform_loop; - VisitStmt(op->body); - - DivergentVarSet next_vars = UnionDivergentVars(loop_entry_vars, divergent_vars_); - UniformBufferSet next_buffers = IntersectUniformBuffers(old_buffers, uniform_buffers_); - if (next_vars == loop_head_vars && next_buffers == loop_head_buffers) { - converged = true; - break; - } - loop_head_vars = std::move(next_vars); - loop_head_buffers = std::move(next_buffers); - } - TVM_FFI_CHECK(converged, ValueError) - << "IKET convergence analysis did not reach a loop fixed point"; - uniform_buffers_ = std::move(loop_head_buffers); - divergent_vars_ = std::move(old_vars); - divergent_context_ = old_divergent; + os << R"IKET( default: + break; } + return event_id; +} - void VisitStmt_(const WhileNode* op) final { - AnnotationFinder annotations; - annotations(op->body); - LoopControlFinder loop_control; - loop_control(op->body); - TVM_FFI_CHECK(!(annotations.found && loop_control.found), ValueError) - << "IKET event sites are not allowed in a loop containing break or continue"; - bool old_divergent = divergent_context_; - DivergentVarSet old_vars = divergent_vars_; - UniformBufferSet old_buffers = uniform_buffers_; - DivergentVarSet loop_head_vars = old_vars; - UniformBufferSet loop_head_buffers = old_buffers; - bool converged = false; - for (size_t iteration = 0; iteration < kMaxConvergenceAnalysisIterations; ++iteration) { - divergent_vars_ = loop_head_vars; - uniform_buffers_ = loop_head_buffers; - VisitExpr(op->condition); - divergent_context_ = old_divergent || !IsUniform(op->condition); - VisitStmt(op->body); - - DivergentVarSet next_vars = UnionDivergentVars(old_vars, divergent_vars_); - UniformBufferSet next_buffers = IntersectUniformBuffers(old_buffers, uniform_buffers_); - if (next_vars == loop_head_vars && next_buffers == loop_head_buffers) { - converged = true; - break; - } - loop_head_vars = std::move(next_vars); - loop_head_buffers = std::move(next_buffers); +template +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id, Payload payload) { + if constexpr (sizeof(Payload) <= 4) { + unsigned int payload_bits = static_cast(payload); + if constexpr (sizeof(Payload) < 4) { + payload_bits &= (1U << (sizeof(Payload) * 8)) - 1U; } - TVM_FFI_CHECK(converged, ValueError) - << "IKET convergence analysis did not reach a loop fixed point"; - uniform_buffers_ = std::move(loop_head_buffers); - divergent_vars_ = std::move(old_vars); - divergent_context_ = old_divergent; + return tvm_builtin_iket_official_event(event_id, payload_bits); + } else { + return tvm_builtin_iket_official_event( + event_id, static_cast(payload)); } +} - void VisitExpr_(const CallNode* call) final { - if (IsIketOp(call->op)) { - bool unproven_token = call->op.same_as(IketRangeEndOp()) && !IsUniform(call->args[0]); - if ((divergent_context_ || unproven_token) && warned_calls_.insert(call).second) { - LOG(WARNING) << "IKET warp convergence could not be proven for event site: " - << GetRef(call) - << "; continuing because convergence diagnostics are advisory"; - } - } - StmtExprVisitor::VisitExpr_(call); - } +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id, float payload) { + return tvm_builtin_iket_official_event(event_id, __float_as_uint(payload)); +} - DivergentVarSet divergent_vars_; - UniformBufferSet uniform_buffers_; - std::unordered_set warned_calls_; - bool divergent_context_{false}; -}; +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id, double payload) { + return tvm_builtin_iket_official_event( + event_id, static_cast(__double_as_longlong(payload))); +} +)IKET"; + return os.str(); +} class InstrumentOfficialKernel : public StmtExprMutator { public: @@ -1429,10 +1121,28 @@ class InstrumentOfficialKernel : public StmtExprMutator { {cast(PrimType::UInt(32), event_id), StringImm(device_source_)}); } + PrimExpr Event(PrimExpr event_id, PrimExpr payload) const { + static const Op& event_op = Op::Get("tirx.cuda.iket_official_event"); + return Call( + PrimType::UInt(32), event_op, + {cast(PrimType::UInt(32), event_id), StringImm(device_source_), std::move(payload)}); + } + + PrimExpr NormalizePayload(PrimExpr payload, PayloadType type) const { + TVM_FFI_ICHECK(type != PayloadType::kNone); + return payload; + } + Stmt VisitStmt_(const EvaluateNode* evaluate) final { if (const auto* call = evaluate->value.as(); call && call->op.same_as(IketRangeEndOp())) { PrimExpr token = VisitExpr(call->args[0]).as_or_throw(); + if (call->args.size() == 2) { + PayloadType payload_type = ValidatePayload(call->args[1]); + PrimExpr payload = + NormalizePayload(VisitExpr(call->args[1]).as_or_throw(), payload_type); + return IfThenElse(token != 0, Evaluate(Event(token, std::move(payload)))); + } return Evaluate(Event(token)); } return StmtExprMutator::VisitStmt_(evaluate); @@ -1440,17 +1150,35 @@ class InstrumentOfficialKernel : public StmtExprMutator { Expr VisitExpr_(const CallNode* call) final { if (call->op.same_as(IketRangeStartOp())) { - return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kRange, call).event_id)); + const Declaration& declaration = Lookup(DeclarationKind::kRange, call); + PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id); + if (declaration.has_payload) { + return Event(event_id, NormalizePayload(VisitExpr(call->args[1]).as_or_throw(), + declaration.payload_type)); + } + return Event(event_id); } if (call->op.same_as(IketSentinelOp())) return IntImm(PrimType::UInt(32), 0); if (call->op.same_as(IketMarkOp())) { - return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kMark, call).event_id)); + const Declaration& declaration = Lookup(DeclarationKind::kMark, call); + PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id); + if (declaration.has_payload) { + return Event(event_id, NormalizePayload(VisitExpr(call->args[1]).as_or_throw(), + declaration.payload_type)); + } + return Event(event_id); } if (call->op.same_as(IketRangePushOp())) { - return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kPush, call).event_id)); + const Declaration& declaration = Lookup(DeclarationKind::kPush, call); + PrimExpr event_id = IntImm(PrimType::UInt(32), declaration.event_id); + if (declaration.has_payload) { + return Event(event_id, NormalizePayload(VisitExpr(call->args[1]).as_or_throw(), + declaration.payload_type)); + } + return Event(event_id); } if (call->op.same_as(IketRangePopOp())) { - return Event(IntImm(PrimType::UInt(32), 31)); + return Event(IntImm(PrimType::UInt(32), kRangePopEventId)); } if (call->op.same_as(IketRangeEndOp())) { TVM_FFI_THROW(ValueError) << "range_end must be emitted in statement position"; @@ -1462,7 +1190,16 @@ class InstrumentOfficialKernel : public StmtExprMutator { std::string device_source_; }; -bool IketEnabled(const IRModule& module) { return module->HasNonzeroAttr("tirx.iket.enabled"); } +bool IketEnabled(const IRModule& module) { + if (module->HasNonzeroAttr("tirx.iket.enabled")) return true; + const char* child_enable = std::getenv("TVM_IKET_INJECTED_CHILD_ENABLE"); + const char* profile = std::getenv("TVM_IKET_OFFICIAL_PROFILE"); + const char* injection = std::getenv("CUDA_INJECTION64_PATH"); + const char* injection_config = std::getenv("SMODEL_INJECTION_CONFIG"); + return child_enable && std::string(child_enable) == "1" && profile && + std::string(profile) == "cutlass-4.6.0" && injection && injection[0] != '\0' && + injection_config && injection_config[0] != '\0'; +} IRModule LowerIketImpl(IRModule module) { if (!IketEnabled(module)) { @@ -1480,6 +1217,7 @@ IRModule LowerIketImpl(IRModule module) { TokenDeclarationMap token_declarations = CollectTokenDeclarations(function->body); RangeEndSchemaVerifier schema_verifier(token_declarations, &collector.declarations); schema_verifier(function->body); + ValidateRangeSchemas(collector.declarations); } StripIket strip(std::move(tokens)); Stmt body = RemoveStrippedIketNoOps()(strip(function->body)); @@ -1503,8 +1241,6 @@ IRModule LowerIketImpl(IRModule module) { std::string function_name = FunctionName(global_var, function); TVM_FFI_CHECK(IsCudaDeviceFunction(function), ValueError) << "IKET annotations are only valid in a split CUDA device kernel"; - TVM_FFI_CHECK_LE(collector.declarations.size(), kMaxDeclarations, ValueError) - << "NVIDIA IKET supports at most " << kMaxDeclarations << " declarations per kernel"; TokenBufferSet tokens = CollectTokenBuffers(function->body); TokenVerifier verifier(tokens); @@ -1512,22 +1248,13 @@ IRModule LowerIketImpl(IRModule module) { TokenDeclarationMap token_declarations = CollectTokenDeclarations(function->body); RangeEndSchemaVerifier schema_verifier(token_declarations, &collector.declarations); schema_verifier(function->body); - IketConvergenceVerifier convergence_verifier; - convergence_verifier(function->body); - + ValidateRangeSchemas(collector.declarations); TVM_FFI_CHECK(IsSm90OrNewer(function), ValueError) << "NVIDIA IKET requires SM90 or newer for kernel " << function_name; - TVM_FFI_CHECK(!HasAnyPayload(collector.declarations), ValueError) - << "NVIDIA IKET does not support payloads in kernel " << function_name; - TVM_FFI_CHECK(OfficialTokenAnalyzer(tokens).Prove(function->body), ValueError) - << "NVIDIA IKET requires token ranges to be provably strictly alternating " - "and closed on every exit in kernel " - << function_name; - TVM_FFI_CHECK(OfficialStackAnalyzer().Prove(function->body), ValueError) - << "NVIDIA IKET requires balanced range_push/range_pop paths in kernel " << function_name; kernels.push_back(KernelIketInfo{global_var, function, std::move(function_name), - std::move(collector.declarations)}); + std::move(collector.declarations), + collector.has_payload_calls}); } std::sort( @@ -1538,29 +1265,50 @@ IRModule LowerIketImpl(IRModule module) { << "IKET device kernels must have unique global symbols: " << kernels[i].name; } - std::map event_ids; + std::map module_declarations; std::unordered_map event_kinds; for (const KernelIketInfo& kernel : kernels) { for (const auto& [key, declaration] : kernel.declarations) { auto [kind_it, kind_inserted] = event_kinds.emplace(key.name, key.kind); TVM_FFI_CHECK(kind_inserted || kind_it->second == key.kind, ValueError) << "NVIDIA IKET declaration " << key.name << " changes event kind across kernels"; - event_ids.emplace(key, 0); + auto [declaration_it, declaration_inserted] = module_declarations.emplace(key, declaration); + if (!declaration_inserted) { + const Declaration& previous = declaration_it->second; + TVM_FFI_CHECK_EQ(previous.has_payload, declaration.has_payload, ValueError) + << "NVIDIA IKET declaration " << key.name << " changes payload presence across kernels"; + TVM_FFI_CHECK(previous.payload_type == declaration.payload_type, TypeError) + << "NVIDIA IKET declaration " << key.name << " changes payload type from " + << PayloadTypeName(previous.payload_type) << " to " + << PayloadTypeName(declaration.payload_type) << " across kernels"; + } } } - TVM_FFI_CHECK_LE(event_ids.size(), kMaxDeclarations, ValueError) - << "NVIDIA IKET supports at most " << kMaxDeclarations - << " distinct declarations in one CUDA module"; + TVM_FFI_CHECK_LE(module_declarations.size(), kExtendedMaxDeclarations, ValueError) + << "NVIDIA IKET supports at most " << kExtendedMaxDeclarations + << " distinct user declarations in one CUDA module; got " << module_declarations.size(); + + InstrumentMode mode = module_declarations.size() <= kNativeMaxDeclarations + ? InstrumentMode::kNativeDump + : InstrumentMode::kExtendedNativeDump; + if (mode == InstrumentMode::kExtendedNativeDump) { + LOG(WARNING) << "NVIDIA IKET is using ExtendedNativeDump for " << module_declarations.size() + << " declarations; records are wider and instrumentation overhead increases"; + } - uint32_t event_id = 1; - for (auto& [key, id] : event_ids) id = event_id++; + uint32_t event_id = + mode == InstrumentMode::kNativeDump ? kNativeFirstEventId : kExtendedFirstEventId; + std::map event_ids; + for (const auto& [key, declaration] : module_declarations) { + event_ids.emplace(key, event_id++); + } for (KernelIketInfo& kernel : kernels) { for (auto& [key, declaration] : kernel.declarations) { declaration.event_id = event_ids.at(key); } } - std::string device_source = BuildOfficialDeviceSource(kernels); + std::string device_source = BuildOfficialDeviceSource(kernels, mode); for (const KernelIketInfo& kernel : kernels) { module->Update(kernel.global_var, InstrumentOfficialKernel(kernel, device_source).Run()); } diff --git a/tests/python/tirx/iket/iket_profile_workload.py b/tests/python/tirx/iket/iket_profile_workload.py index 9e2218ba84a8..87c81dd24bdd 100644 --- a/tests/python/tirx/iket/iket_profile_workload.py +++ b/tests/python/tirx/iket/iket_profile_workload.py @@ -50,6 +50,66 @@ def canonical_iket_workload(out: T.Buffer((32,), "int32")): out[tx] = tx + 1 +@T.prim_func +def native_payload_workload(out: T.Buffer((32,), "int32")): + T.device_entry() + profiler = iket.IketProfiler() + tx = T.thread_id([32]) + profiler.mark("lane_payload", tx + 100) + if tx >= 5: + profiler.mark("first_active_lane", tx) + profiler.mark("wide_payload", T.int64(tx) + T.int64(0x100000000)) + profiler.mark("negative_payload", T.int32(-32)) + profiler.mark("bool_true_payload", tx == 0) + profiler.mark("bool_false_payload", tx != 0) + profiler.mark("float32_payload", T.float32(-3.25)) + profiler.mark("float64_payload", T.float64(6.5)) + token = profiler.range_start("token_payload", tx + 200) + profiler.range_end(token, tx + 300) + profiler.range_push("stack_payload", tx + 400) + profiler.range_pop() + out[tx] = tx + 2 + + +@T.prim_func +def extended_payload_workload(out: T.Buffer((32,), "int32")): + T.device_entry() + profiler = iket.IketProfiler() + tx = T.thread_id([32]) + profiler.mark("extended_lane_payload", tx + 500) + profiler.mark("extended01") + profiler.mark("extended02") + profiler.mark("extended03") + profiler.mark("extended04") + profiler.mark("extended05") + profiler.mark("extended06") + profiler.mark("extended07") + profiler.mark("extended08") + profiler.mark("extended09") + profiler.mark("extended10") + profiler.mark("extended11") + profiler.mark("extended12") + profiler.mark("extended13") + profiler.mark("extended14") + profiler.mark("extended15") + profiler.mark("extended16") + profiler.mark("extended17") + profiler.mark("extended18") + profiler.mark("extended19") + profiler.mark("extended20") + profiler.mark("extended21") + profiler.mark("extended22") + profiler.mark("extended23") + profiler.mark("extended24") + profiler.mark("extended25") + profiler.mark("extended26") + profiler.mark("extended27") + profiler.mark("extended28") + profiler.mark("extended29") + profiler.mark("extended30") + out[tx] = tx + 3 + + def _parse_args(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output-dir", default="/tmp/tvm-iket-workload") @@ -80,15 +140,24 @@ def _injection_tool_name(): def _profile_workload(args): if args.fail_capture and _injection_tool_name() == "iket": raise RuntimeError("intentional capture-only IKET workload failure") - executable = iket.IketProfiler().compile( - canonical_iket_workload, - target=tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}), - tir_pipeline="tirx", + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + workloads = ( + (canonical_iket_workload, 1), + (native_payload_workload, 2), + (extended_payload_workload, 3), ) - out = tvm.runtime.empty((32,), "int32", device=tvm.cuda()) - executable["canonical_iket_workload"](out) + outputs = [] + for workload, offset in workloads: + # Plain JIT compilation is intentionally used here. The validated + # run-iket child enables LowerIket automatically for these modules. + executable = tvm.compile(workload, target=target, tir_pipeline="tirx") + module = executable.jit() + out = tvm.runtime.empty((32,), "int32", device=tvm.cuda()) + module.main(out) + outputs.append((out, offset)) tvm.cuda().sync() - np.testing.assert_array_equal(out.numpy(), np.arange(32, dtype=np.int32) + 1) + for out, offset in outputs: + np.testing.assert_array_equal(out.numpy(), np.arange(32, dtype=np.int32) + offset) def main(): diff --git a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py index 2b4a99256477..619090a114f0 100644 --- a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py +++ b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py @@ -6,7 +6,7 @@ # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -36,8 +36,8 @@ @cute.kernel -def oracle_kernel(): - """Exercise every no-payload operation supported by the official backend.""" +def native_no_payload_kernel(): + """Exercise every no-payload operation in NativeDump mode.""" cute_iket.mark("mark") token = cute_iket.range_start("token") cute_iket.range_end(token) @@ -45,11 +45,162 @@ def oracle_kernel(): cute_iket.range_pop() cute_iket.mark("slash/name") cute_iket.mark("é") + _native_no_payload_fillers() @cute.jit -def oracle_launch(): - oracle_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) +def native_no_payload_launch(): + native_no_payload_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) + + +def _payload_events(): + """Emit the public scalar payload matrix and both range forms.""" + # CUTLASS 4.6.0's public helper currently forwards Python bool as i1 even + # though the IKET dialect requires an integer payload of at least 8 bits. + # UI8 is the verified wire representation used by the TIRx bool shim. + cute_iket.mark("bool", cutlass.Uint8(True)) + cute_iket.mark("i8", cutlass.Int8(-8)) + cute_iket.mark("u8", cutlass.Uint8(8)) + cute_iket.mark("i16", cutlass.Int16(-16)) + cute_iket.mark("u16", cutlass.Uint16(16)) + cute_iket.mark("i32", cutlass.Int32(-32)) + cute_iket.mark("u32", cutlass.Uint32(32)) + cute_iket.mark("i64", cutlass.Int64(-64)) + cute_iket.mark("u64", cutlass.Uint64(64)) + cute_iket.mark("f32", cutlass.Float32(-3.25)) + cute_iket.mark("f64", cutlass.Float64(6.5)) + token = cute_iket.range_start("token_payload", cutlass.Int32(-7)) + cute_iket.range_end(token, cutlass.Int32(9)) + cute_iket.range_push("stack_payload", cutlass.Float32(1.5)) + cute_iket.range_pop() + + +def _native_no_payload_fillers(): + """Bring the five operation-covering declarations to the Native limit.""" + cute_iket.mark("native_filler00") + cute_iket.mark("native_filler01") + cute_iket.mark("native_filler02") + cute_iket.mark("native_filler03") + cute_iket.mark("native_filler04") + cute_iket.mark("native_filler05") + cute_iket.mark("native_filler06") + cute_iket.mark("native_filler07") + cute_iket.mark("native_filler08") + cute_iket.mark("native_filler09") + cute_iket.mark("native_filler10") + cute_iket.mark("native_filler11") + cute_iket.mark("native_filler12") + cute_iket.mark("native_filler13") + cute_iket.mark("native_filler14") + cute_iket.mark("native_filler15") + cute_iket.mark("native_filler16") + cute_iket.mark("native_filler17") + cute_iket.mark("native_filler18") + cute_iket.mark("native_filler19") + cute_iket.mark("native_filler20") + cute_iket.mark("native_filler21") + cute_iket.mark("native_filler22") + cute_iket.mark("native_filler23") + cute_iket.mark("native_filler24") + + +def _native_payload_fillers(): + """Bring the thirteen payload declarations to the Native limit.""" + cute_iket.mark("filler00") + cute_iket.mark("filler01") + cute_iket.mark("filler02") + cute_iket.mark("filler03") + cute_iket.mark("filler04") + cute_iket.mark("filler05") + cute_iket.mark("filler06") + cute_iket.mark("filler07") + cute_iket.mark("filler08") + cute_iket.mark("filler09") + cute_iket.mark("filler10") + cute_iket.mark("filler11") + cute_iket.mark("filler12") + cute_iket.mark("filler13") + cute_iket.mark("filler14") + cute_iket.mark("filler15") + cute_iket.mark("filler16") + + +@cute.kernel +def native_payload_kernel(): + """Exercise all public payload types in NativeDump mode.""" + _payload_events() + _native_payload_fillers() + + +@cute.jit +def native_payload_launch(): + native_payload_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) + + +def _extended_fillers(): + _native_payload_fillers() + cute_iket.mark("filler17") + + +@cute.kernel +def extended_no_payload_kernel(): + """Use 31 distinct names to select ExtendedNativeDump.""" + cute_iket.mark("event00") + cute_iket.mark("event01") + cute_iket.mark("event02") + cute_iket.mark("event03") + cute_iket.mark("event04") + cute_iket.mark("event05") + cute_iket.mark("event06") + cute_iket.mark("event07") + cute_iket.mark("event08") + cute_iket.mark("event09") + cute_iket.mark("event10") + cute_iket.mark("event11") + cute_iket.mark("event12") + cute_iket.mark("event13") + cute_iket.mark("event14") + cute_iket.mark("event15") + cute_iket.mark("event16") + cute_iket.mark("event17") + cute_iket.mark("event18") + cute_iket.mark("event19") + cute_iket.mark("event20") + cute_iket.mark("event21") + cute_iket.mark("event22") + cute_iket.mark("event23") + cute_iket.mark("event24") + cute_iket.mark("event25") + cute_iket.mark("event26") + cute_iket.mark("event27") + cute_iket.mark("event28") + cute_iket.mark("event29") + cute_iket.mark("event30") + + +@cute.jit +def extended_no_payload_launch(): + extended_no_payload_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) + + +@cute.kernel +def extended_payload_kernel(): + """Exercise payloads with 31 names in ExtendedNativeDump mode.""" + _payload_events() + _extended_fillers() + + +@cute.jit +def extended_payload_launch(): + extended_payload_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) + + +CASES = { + "native_no_payload": native_no_payload_launch, + "native_payload": native_payload_launch, + "extended_no_payload": extended_no_payload_launch, + "extended_payload": extended_payload_launch, +} def _normalize(value): @@ -90,6 +241,14 @@ def _command_output(command): return subprocess.run(command, check=True, capture_output=True, text=True).stdout.strip() +def _packaged_nvdisasm() -> Path: + distribution = metadata.distribution("nvidia-cuda-nvdisasm") + path = Path(distribution.locate_file("nvidia/cu13/bin/nvdisasm")) + if not path.is_file(): + raise RuntimeError(f"packaged nvdisasm is missing: {path}") + return path + + def _normalized_metadata(context): events = sorted(_normalize(context.get_all_events()), key=lambda item: item["event_id"]) ranges = sorted(_normalize(context.get_all_ranges()), key=lambda item: item["range_id"]) @@ -109,51 +268,98 @@ def _normalized_metadata(context): } +def _instruction_features(ptx: str, sass: str) -> dict: + ptx_features = ( + "activemask", + "elect.sync", + "globaltimer_lo", + "st.weak.shared.u32", + "st.weak.shared.u64", + "pmevent.mask", + ) + sass_features = ( + "S2R", + "ELECT", + "CS2R", + "STS", + "PMTRIG", + ) + return { + "ptx": {feature: ptx.count(feature) for feature in ptx_features}, + "sass": {feature: sass.count(feature) for feature in sass_features}, + } + + +def _compile_case(case_name: str, launch, output_dir: Path, nvdisasm: Path) -> dict: + case_dir = output_dir / case_name + case_dir.mkdir(parents=True, exist_ok=True) + compiled = cute.compile( + launch, + options=(f"iket --dump-dir={case_dir} --keep-ptx --keep-cubin --keep-sass"), + ) + ptx = _artifact_bytes(compiled.artifacts.PTX) + cubin = _artifact_bytes(compiled.artifacts.CUBIN) + cubin_path = case_dir / f"{case_name}.cubin" + cubin_path.write_bytes(cubin) + sass = subprocess.run( + [nvdisasm, "-c", cubin_path], check=True, capture_output=True, text=True + ).stdout + context = iket.Context(cubin) + if not context.is_instrumented(): + raise RuntimeError(f"CUTLASS DSL oracle case {case_name} is not IKET-instrumented") + + normalized_metadata = _normalized_metadata(context) + methods = sorted( + { + event["instrument_method"] + for event in normalized_metadata["events"] + if event["event_id"] not in (0, 31) + } + ) + return { + "instrument_methods": methods, + "artifact_sha256": { + "ptx": _sha256(ptx), + "cubin": _sha256(cubin), + "sass": _sha256(sass), + }, + "instruction_features": _instruction_features(ptx.decode(), sass), + "metadata": normalized_metadata, + "metadata_sha256": _sha256( + json.dumps(normalized_metadata, sort_keys=True, separators=(",", ":")) + ), + } + + +def _wheel_hashes(wheels: list[Path]) -> dict[str, str]: + if wheels: + return {wheel.name: _sha256(wheel) for wheel in sorted(wheels)} + prior_manifest = Path(__file__).with_name("iket_official_cutlass_4_6_0_oracle.json") + if prior_manifest.is_file(): + prior = json.loads(prior_manifest.read_text(encoding="utf-8")) + if prior.get("profile", {}).get("cutlass_dsl") == "4.6.0": + return prior.get("wheels", {}) + return {} + + def main(): parser = argparse.ArgumentParser() parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--manifest-path", type=Path) parser.add_argument("--wheel", type=Path, action="append", default=[]) - for artifact_name in ( - "unpatched-kernel", - "patched-kernel", - "unpatched-disassembly", - "patched-disassembly", - ): - parser.add_argument(f"--{artifact_name}", type=Path) args = parser.parse_args() - patch_artifacts = { - name: getattr(args, name) - for name in ( - "unpatched_kernel", - "patched_kernel", - "unpatched_disassembly", - "patched_disassembly", - ) - } - if any(patch_artifacts.values()) and not all(patch_artifacts.values()): - parser.error("all four patched/unpatched artifact paths must be supplied together") args.output_dir.mkdir(parents=True, exist_ok=True) cutlass.cuda.initialize_cuda_context() - compiled = cute.compile( - oracle_launch, - options=(f"iket --dump-dir={args.output_dir} --keep-ptx --keep-cubin --keep-sass"), - ) - artifacts = { - name.lower(): _artifact_bytes(getattr(compiled.artifacts, name)) - for name in ("PTX", "CUBIN", "SASS") + nvdisasm = _packaged_nvdisasm() + cases = { + name: _compile_case(name, launch, args.output_dir, nvdisasm) + for name, launch in CASES.items() } - cubin = artifacts["cubin"] - context = iket.Context(cubin) - if not context.is_instrumented(): - raise RuntimeError("CUTLASS DSL oracle CUBIN is not IKET-instrumented") - - normalized_metadata = _normalized_metadata(context) manifest = { - "schema_version": 2, + "schema_version": 3, "profile": { "cutlass_dsl": metadata.version("nvidia-cutlass-dsl"), - "instrument_method": "NativeDump", "driver": _command_output( [ "nvidia-smi", @@ -162,48 +368,46 @@ def main(): "--id=0", ] ), - "nvdisasm": _command_output(["nvdisasm", "--version"]), + "nvdisasm": _command_output([nvdisasm, "--version"]), + "nvdisasm_distribution": metadata.version("nvidia-cuda-nvdisasm"), + "nvrtc_distribution": metadata.version("nvidia-cuda-nvrtc"), "compiler_flags": [ "iket", - "--dump-dir=", + "--dump-dir=/", "--keep-ptx", "--keep-cubin", "--keep-sass", ], }, - "wheels": {wheel.name: _sha256(wheel) for wheel in sorted(args.wheel)}, - "artifact_sha256": {name: _sha256(value) for name, value in artifacts.items()}, - "metadata": normalized_metadata, - "metadata_sha256": _sha256( - json.dumps(normalized_metadata, sort_keys=True, separators=(",", ":")) - ), - "native_dump_abi": { + "wheels": _wheel_hashes(args.wheel), + "cases": cases, + "abi": { "meta_info_bytes": 48, "event_attributes_bytes": 60, "range_attributes_bytes": 72, "sentinel_event_id": 0, "range_pop_event_id": 31, - "placeholder": [ - "READ_CLUSTER_CTARANK", - "GLOBALTIMERLO", - "ENCODE_EVENT_ID", - "COMPUTE_SHARED_PLACEHOLDER_ADDRESS", - "STORE_SHARED_WEAK_32", - "PMEVENT", - ], - "patched_hot_path": [ + "native_user_event_ids": [1, 30], + "extended_user_event_ids": [64, 4095], + "max_user_declarations": 4032, + "native_patched_hot_path": [ "GLOBALTIMERLO", "ENCODE_EVENT_ID", "STORE_GLOBAL_32", "ADD_WRITE_PTR_64_4", ], + "record_layouts": { + "native_no_payload": {"timestamp_event": [0, 4]}, + "native_payload_32": {"timestamp_event": [0, 4], "payload": [4, 4]}, + "native_payload_64": {"timestamp_event": [0, 4], "payload": [8, 8]}, + "extended_no_payload": {"timestamp_event": [0, 8]}, + "extended_payload_32": {"timestamp_event": [0, 8], "payload": [8, 4]}, + "extended_payload_64": {"timestamp_event": [0, 8], "payload": [8, 8]}, + }, }, } - if all(patch_artifacts.values()): - manifest["patch_artifact_sha256"] = { - name: _sha256(path) for name, path in patch_artifacts.items() - } - manifest_path = args.output_dir / "iket_official_oracle.json" + manifest_path = args.manifest_path or args.output_dir / "iket_official_oracle.json" + manifest_path.parent.mkdir(parents=True, exist_ok=True) manifest_path.write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" ) diff --git a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json index 0bdb9e50f142..ecbe88720325 100644 --- a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json +++ b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json @@ -1,111 +1,1511 @@ { - "artifact_sha256": { - "cubin": "ddd0cbfc511914d453cfc9ba770bddc3939a6ce77f13488eae1aa31bd867a1f5", - "ptx": "9a679b1ceefe35bd4f783d13bcf7b8a4c1490f94b606323cd9ffb7b1c5aa7554", - "sass": "5d28eac4ad110174ab8fb99d1437ebf0e99b4de4df6c1cee2341e2ef210f5af7" - }, - "metadata": { - "events": [ - { - "event_id": 0, - "event_name": "reserved_event_id_0", - "event_pos": 0, - "instrument_method": 3, - "payload_type": 0, - "range_id": 0 - }, - { - "event_id": 1, - "event_name": "mark", - "event_pos": 0, - "instrument_method": 3, - "payload_type": 0, - "range_id": 0 - }, - { - "event_id": 2, - "event_name": "token", - "event_pos": 4, - "instrument_method": 3, - "payload_type": 0, - "range_id": 2491017778 - }, - { - "event_id": 3, - "event_name": "stack", - "event_pos": 1, - "instrument_method": 3, - "payload_type": 0, - "range_id": 1649501183 - }, - { - "event_id": 4, - "event_name": "slash/name", - "event_pos": 0, - "instrument_method": 3, - "payload_type": 0, - "range_id": 0 - }, - { - "event_id": 5, - "event_name": "\u00e9", - "event_pos": 0, - "instrument_method": 3, - "payload_type": 0, - "range_id": 0 - }, - { - "event_id": 31, - "event_name": "pop_range", - "event_pos": 2, - "instrument_method": 3, - "payload_type": 0, - "range_id": 0 - } + "abi": { + "event_attributes_bytes": 60, + "extended_user_event_ids": [ + 64, + 4095 + ], + "max_user_declarations": 4032, + "meta_info_bytes": 48, + "native_patched_hot_path": [ + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "STORE_GLOBAL_32", + "ADD_WRITE_PTR_64_4" + ], + "native_user_event_ids": [ + 1, + 30 ], - "info": { - "legacy_iket_cubin": 0, - "magic_number": 3133075869, - "max_event_id": 31, - "max_event_name_size": 32, - "supported_features": 3, - "version_major": 0, - "version_minor": 5 + "range_attributes_bytes": 72, + "range_pop_event_id": 31, + "record_layouts": { + "extended_no_payload": { + "timestamp_event": [ + 0, + 8 + ] + }, + "extended_payload_32": { + "payload": [ + 8, + 4 + ], + "timestamp_event": [ + 0, + 8 + ] + }, + "extended_payload_64": { + "payload": [ + 8, + 8 + ], + "timestamp_event": [ + 0, + 8 + ] + }, + "native_no_payload": { + "timestamp_event": [ + 0, + 4 + ] + }, + "native_payload_32": { + "payload": [ + 4, + 4 + ], + "timestamp_event": [ + 0, + 4 + ] + }, + "native_payload_64": { + "payload": [ + 8, + 8 + ], + "timestamp_event": [ + 0, + 4 + ] + } + }, + "sentinel_event_id": 0 + }, + "cases": { + "extended_no_payload": { + "artifact_sha256": { + "cubin": "2efbee48a8fed94b2b05beac11456ea3fc5a14dc6737325853dda8e59bd0b72d", + "ptx": "5978f1a3a5ce3c7b56d0dba1531e1bb341a2c43deda8fbf1e0112fcfd7d9a36e", + "sass": "46a07c67fa0f8d2491884cb26fd6774b63a50811e3db020dace3e1fc25e4c6b0" + }, + "instruction_features": { + "ptx": { + "activemask": 0, + "elect.sync": 0, + "globaltimer_lo": 31, + "pmevent.mask": 31, + "st.weak.shared.u32": 0, + "st.weak.shared.u64": 31 + }, + "sass": { + "CS2R": 0, + "ELECT": 0, + "PMTRIG": 31, + "S2R": 0, + "STS": 31 + } + }, + "instrument_methods": [ + 5 + ], + "metadata": { + "events": [ + { + "event_id": 0, + "event_name": "reserved_event_id_0", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 64, + "event_name": "event00", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 65, + "event_name": "event01", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 66, + "event_name": "event02", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 67, + "event_name": "event03", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 68, + "event_name": "event04", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 69, + "event_name": "event05", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 70, + "event_name": "event06", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 71, + "event_name": "event07", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 72, + "event_name": "event08", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 73, + "event_name": "event09", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 74, + "event_name": "event10", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 75, + "event_name": "event11", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 76, + "event_name": "event12", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 77, + "event_name": "event13", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 78, + "event_name": "event14", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 79, + "event_name": "event15", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 80, + "event_name": "event16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 81, + "event_name": "event17", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 82, + "event_name": "event18", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 83, + "event_name": "event19", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 84, + "event_name": "event20", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 85, + "event_name": "event21", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 86, + "event_name": "event22", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 87, + "event_name": "event23", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 88, + "event_name": "event24", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 89, + "event_name": "event25", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 90, + "event_name": "event26", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 91, + "event_name": "event27", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 92, + "event_name": "event28", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 93, + "event_name": "event29", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 94, + "event_name": "event30", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + } + ], + "info": { + "legacy_iket_cubin": 0, + "magic_number": 3133075869, + "max_event_id": 4095, + "max_event_name_size": 32, + "supported_features": 3, + "version_major": 0, + "version_minor": 5 + }, + "kernels": [ + { + "event_sequence": [ + { + "event_id": 64, + "event_name": "event00", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 65, + "event_name": "event01", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 66, + "event_name": "event02", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 67, + "event_name": "event03", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 68, + "event_name": "event04", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 69, + "event_name": "event05", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 70, + "event_name": "event06", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 71, + "event_name": "event07", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 72, + "event_name": "event08", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 73, + "event_name": "event09", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 74, + "event_name": "event10", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 75, + "event_name": "event11", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 76, + "event_name": "event12", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 77, + "event_name": "event13", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 78, + "event_name": "event14", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 79, + "event_name": "event15", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 80, + "event_name": "event16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 81, + "event_name": "event17", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 82, + "event_name": "event18", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 83, + "event_name": "event19", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 84, + "event_name": "event20", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 85, + "event_name": "event21", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 86, + "event_name": "event22", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 87, + "event_name": "event23", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 88, + "event_name": "event24", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 89, + "event_name": "event25", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 90, + "event_name": "event26", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 91, + "event_name": "event27", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 92, + "event_name": "event28", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 93, + "event_name": "event29", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 94, + "event_name": "event30", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + } + ], + "kernel_name": "kernel_cutlass_extended_no_payload_kernel_0" + } + ], + "ranges": [] + }, + "metadata_sha256": "70e8b770eee1c4f80e86c88df44d27df20c4805f68ba3a51289f30cb897bfdc2" }, - "kernels": [ - { - "event_sequence": [ + "extended_payload": { + "artifact_sha256": { + "cubin": "548598da0def27041297295c37696c35d2027085577dc167e4125a00d55daafd", + "ptx": "b8929715c2dd98e55bd1227f90d3b8a333643fb13d65e56225cda5a83291bef8", + "sass": "9bf3690438505e9fb0e1d5241aaca819b5eef2446dd627e925c5b8da778edeaa" + }, + "instruction_features": { + "ptx": { + "activemask": 14, + "elect.sync": 14, + "globaltimer_lo": 33, + "pmevent.mask": 33, + "st.weak.shared.u32": 0, + "st.weak.shared.u64": 33 + }, + "sass": { + "CS2R": 0, + "ELECT": 14, + "PMTRIG": 33, + "S2R": 0, + "STS": 47 + } + }, + "instrument_methods": [ + 5 + ], + "metadata": { + "events": [ + { + "event_id": 0, + "event_name": "reserved_event_id_0", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 64, + "event_name": "bool", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 65, + "event_name": "i8", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 1, + "range_id": 0 + }, + { + "event_id": 66, + "event_name": "u8", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 67, + "event_name": "i16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 3, + "range_id": 0 + }, + { + "event_id": 68, + "event_name": "u16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 4, + "range_id": 0 + }, + { + "event_id": 69, + "event_name": "i32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 5, + "range_id": 0 + }, + { + "event_id": 70, + "event_name": "u32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 6, + "range_id": 0 + }, + { + "event_id": 71, + "event_name": "i64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 7, + "range_id": 0 + }, + { + "event_id": 72, + "event_name": "u64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 16, + "range_id": 0 + }, + { + "event_id": 73, + "event_name": "f32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 13, + "range_id": 0 + }, + { + "event_id": 74, + "event_name": "f64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 14, + "range_id": 0 + }, + { + "event_id": 75, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 5, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 76, + "event_name": "stack_payload", + "event_pos": 1, + "instrument_method": 5, + "payload_type": 13, + "range_id": 3918402126 + }, + { + "event_id": 77, + "event_name": "filler00", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 78, + "event_name": "filler01", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 79, + "event_name": "filler02", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 80, + "event_name": "filler03", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 81, + "event_name": "filler04", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 82, + "event_name": "filler05", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 83, + "event_name": "filler06", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 84, + "event_name": "filler07", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 85, + "event_name": "filler08", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 86, + "event_name": "filler09", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 87, + "event_name": "filler10", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 88, + "event_name": "filler11", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 89, + "event_name": "filler12", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 90, + "event_name": "filler13", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 91, + "event_name": "filler14", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 92, + "event_name": "filler15", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 93, + "event_name": "filler16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 94, + "event_name": "filler17", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + } + ], + "info": { + "legacy_iket_cubin": 0, + "magic_number": 3133075869, + "max_event_id": 4095, + "max_event_name_size": 32, + "supported_features": 3, + "version_major": 0, + "version_minor": 5 + }, + "kernels": [ + { + "event_sequence": [ + { + "event_id": 64, + "event_name": "bool", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 65, + "event_name": "i8", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 1, + "range_id": 0 + }, + { + "event_id": 66, + "event_name": "u8", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 67, + "event_name": "i16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 3, + "range_id": 0 + }, + { + "event_id": 68, + "event_name": "u16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 4, + "range_id": 0 + }, + { + "event_id": 69, + "event_name": "i32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 5, + "range_id": 0 + }, + { + "event_id": 70, + "event_name": "u32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 6, + "range_id": 0 + }, + { + "event_id": 71, + "event_name": "i64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 7, + "range_id": 0 + }, + { + "event_id": 72, + "event_name": "u64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 16, + "range_id": 0 + }, + { + "event_id": 73, + "event_name": "f32", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 13, + "range_id": 0 + }, + { + "event_id": 74, + "event_name": "f64", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 14, + "range_id": 0 + }, + { + "event_id": 75, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 5, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 75, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 5, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 76, + "event_name": "stack_payload", + "event_pos": 1, + "instrument_method": 5, + "payload_type": 13, + "range_id": 3918402126 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 77, + "event_name": "filler00", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 78, + "event_name": "filler01", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 79, + "event_name": "filler02", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 80, + "event_name": "filler03", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 81, + "event_name": "filler04", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 82, + "event_name": "filler05", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 83, + "event_name": "filler06", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 84, + "event_name": "filler07", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 85, + "event_name": "filler08", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 86, + "event_name": "filler09", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 87, + "event_name": "filler10", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 88, + "event_name": "filler11", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 89, + "event_name": "filler12", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 90, + "event_name": "filler13", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 91, + "event_name": "filler14", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 92, + "event_name": "filler15", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 93, + "event_name": "filler16", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 94, + "event_name": "filler17", + "event_pos": 0, + "instrument_method": 5, + "payload_type": 0, + "range_id": 0 + } + ], + "kernel_name": "kernel_cutlass_extended_payload_kernel_0" + } + ], + "ranges": [ + { + "color": -1, + "evt_pair_mode": 1, + "range_id": 742653351, + "range_name": "token_payload", + "range_scope": 0, + "range_type": 1 + }, + { + "color": -1, + "evt_pair_mode": 0, + "range_id": 3918402126, + "range_name": "stack_payload", + "range_scope": 0, + "range_type": 2 + } + ] + }, + "metadata_sha256": "4a5b992d31f0c7e5c3f8e55c30f587354499d2ba4c3b6f76d1ce2299196d8fdc" + }, + "native_no_payload": { + "artifact_sha256": { + "cubin": "1cce96680a5cbf97636095401845d23352cf76c1c3780e2ceceb23f04f06c457", + "ptx": "9ee395a247733937e0bea5aff1909c7651de674aa32784147339bd0a8cb4a351", + "sass": "efd375267ba06d6a9b6e222493259895437578878a7f420b680e06f07dcc1841" + }, + "instruction_features": { + "ptx": { + "activemask": 0, + "elect.sync": 0, + "globaltimer_lo": 32, + "pmevent.mask": 32, + "st.weak.shared.u32": 32, + "st.weak.shared.u64": 0 + }, + "sass": { + "CS2R": 32, + "ELECT": 0, + "PMTRIG": 32, + "S2R": 32, + "STS": 32 + } + }, + "instrument_methods": [ + 3 + ], + "metadata": { + "events": [ + { + "event_id": 0, + "event_name": "reserved_event_id_0", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 1, + "event_name": "mark", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 3, + "event_name": "stack", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 0, + "range_id": 1649501183 + }, + { + "event_id": 4, + "event_name": "slash/name", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 5, + "event_name": "\u00e9", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 6, + "event_name": "native_filler00", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 7, + "event_name": "native_filler01", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 8, + "event_name": "native_filler02", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 9, + "event_name": "native_filler03", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 10, + "event_name": "native_filler04", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 11, + "event_name": "native_filler05", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 12, + "event_name": "native_filler06", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 13, + "event_name": "native_filler07", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 14, + "event_name": "native_filler08", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 15, + "event_name": "native_filler09", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 16, + "event_name": "native_filler10", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 17, + "event_name": "native_filler11", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 18, + "event_name": "native_filler12", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 19, + "event_name": "native_filler13", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 20, + "event_name": "native_filler14", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 21, + "event_name": "native_filler15", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 22, + "event_name": "native_filler16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, { - "event_id": 1, - "event_name": "mark", + "event_id": 23, + "event_name": "native_filler17", "event_pos": 0, "instrument_method": 3, "payload_type": 0, "range_id": 0 }, { - "event_id": 2, - "event_name": "token", - "event_pos": 4, + "event_id": 24, + "event_name": "native_filler18", + "event_pos": 0, "instrument_method": 3, "payload_type": 0, - "range_id": 2491017778 + "range_id": 0 }, { - "event_id": 2, - "event_name": "token", - "event_pos": 4, + "event_id": 25, + "event_name": "native_filler19", + "event_pos": 0, "instrument_method": 3, "payload_type": 0, - "range_id": 2491017778 + "range_id": 0 }, { - "event_id": 3, - "event_name": "stack", - "event_pos": 1, + "event_id": 26, + "event_name": "native_filler20", + "event_pos": 0, "instrument_method": 3, "payload_type": 0, - "range_id": 1649501183 + "range_id": 0 + }, + { + "event_id": 27, + "event_name": "native_filler21", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 28, + "event_name": "native_filler22", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 29, + "event_name": "native_filler23", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 30, + "event_name": "native_filler24", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 }, { "event_id": 31, @@ -114,88 +1514,895 @@ "instrument_method": 3, "payload_type": 0, "range_id": 0 + } + ], + "info": { + "legacy_iket_cubin": 0, + "magic_number": 3133075869, + "max_event_id": 31, + "max_event_name_size": 32, + "supported_features": 3, + "version_major": 0, + "version_minor": 5 + }, + "kernels": [ + { + "event_sequence": [ + { + "event_id": 1, + "event_name": "mark", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 3, + "event_name": "stack", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 0, + "range_id": 1649501183 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 4, + "event_name": "slash/name", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 5, + "event_name": "\u00e9", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 6, + "event_name": "native_filler00", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 7, + "event_name": "native_filler01", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 8, + "event_name": "native_filler02", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 9, + "event_name": "native_filler03", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 10, + "event_name": "native_filler04", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 11, + "event_name": "native_filler05", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 12, + "event_name": "native_filler06", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 13, + "event_name": "native_filler07", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 14, + "event_name": "native_filler08", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 15, + "event_name": "native_filler09", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 16, + "event_name": "native_filler10", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 17, + "event_name": "native_filler11", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 18, + "event_name": "native_filler12", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 19, + "event_name": "native_filler13", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 20, + "event_name": "native_filler14", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 21, + "event_name": "native_filler15", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 22, + "event_name": "native_filler16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 23, + "event_name": "native_filler17", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 24, + "event_name": "native_filler18", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 25, + "event_name": "native_filler19", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 26, + "event_name": "native_filler20", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 27, + "event_name": "native_filler21", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 28, + "event_name": "native_filler22", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 29, + "event_name": "native_filler23", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 30, + "event_name": "native_filler24", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + } + ], + "kernel_name": "kernel_cutlass_native_no_payload_kernel_0" + } + ], + "ranges": [ + { + "color": -1, + "evt_pair_mode": 0, + "range_id": 1649501183, + "range_name": "stack", + "range_scope": 0, + "range_type": 2 }, { - "event_id": 4, - "event_name": "slash/name", + "color": -1, + "evt_pair_mode": 1, + "range_id": 2491017778, + "range_name": "token", + "range_scope": 0, + "range_type": 1 + } + ] + }, + "metadata_sha256": "6bf57629fdc493af5b5b660c335d4d6b21465a4fe01ebd3042ed252bf220a071" + }, + "native_payload": { + "artifact_sha256": { + "cubin": "8c3ecdab58bb9f899b58e876d0405af024292ad45dd45b2e18a26ab2588d7caa", + "ptx": "b7a5c0cc80b38a7517ae89e9dec24de927df855fd3963d261992e8d8213e5c69", + "sass": "2bf58bebfeeda55652c9e6db78fba7c391124605100c2d1229364198088ecbd0" + }, + "instruction_features": { + "ptx": { + "activemask": 14, + "elect.sync": 14, + "globaltimer_lo": 32, + "pmevent.mask": 32, + "st.weak.shared.u32": 32, + "st.weak.shared.u64": 0 + }, + "sass": { + "CS2R": 32, + "ELECT": 14, + "PMTRIG": 32, + "S2R": 32, + "STS": 46 + } + }, + "instrument_methods": [ + 3 + ], + "metadata": { + "events": [ + { + "event_id": 0, + "event_name": "reserved_event_id_0", "event_pos": 0, "instrument_method": 3, "payload_type": 0, "range_id": 0 }, + { + "event_id": 1, + "event_name": "bool", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "i8", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 1, + "range_id": 0 + }, + { + "event_id": 3, + "event_name": "u8", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 4, + "event_name": "i16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 3, + "range_id": 0 + }, { "event_id": 5, - "event_name": "\u00e9", + "event_name": "u16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 4, + "range_id": 0 + }, + { + "event_id": 6, + "event_name": "i32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 5, + "range_id": 0 + }, + { + "event_id": 7, + "event_name": "u32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 6, + "range_id": 0 + }, + { + "event_id": 8, + "event_name": "i64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 7, + "range_id": 0 + }, + { + "event_id": 9, + "event_name": "u64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 16, + "range_id": 0 + }, + { + "event_id": 10, + "event_name": "f32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 13, + "range_id": 0 + }, + { + "event_id": 11, + "event_name": "f64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 14, + "range_id": 0 + }, + { + "event_id": 12, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 13, + "event_name": "stack_payload", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 13, + "range_id": 3918402126 + }, + { + "event_id": 14, + "event_name": "filler00", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 15, + "event_name": "filler01", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 16, + "event_name": "filler02", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 17, + "event_name": "filler03", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 18, + "event_name": "filler04", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 19, + "event_name": "filler05", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 20, + "event_name": "filler06", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 21, + "event_name": "filler07", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 22, + "event_name": "filler08", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 23, + "event_name": "filler09", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 24, + "event_name": "filler10", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 25, + "event_name": "filler11", "event_pos": 0, "instrument_method": 3, "payload_type": 0, "range_id": 0 + }, + { + "event_id": 26, + "event_name": "filler12", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 27, + "event_name": "filler13", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 28, + "event_name": "filler14", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 29, + "event_name": "filler15", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 30, + "event_name": "filler16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 } ], - "kernel_name": "kernel_cutlass_oracle_kernel_0" - } - ], - "ranges": [ - { - "color": -1, - "evt_pair_mode": 0, - "range_id": 1649501183, - "range_name": "stack", - "range_scope": 0, - "range_type": 2 - }, - { - "color": -1, - "evt_pair_mode": 1, - "range_id": 2491017778, - "range_name": "token", - "range_scope": 0, - "range_type": 1 - } - ] - }, - "metadata_sha256": "85c13d2584ca6a37b6bfd17dc17d70ded6696f63fae728cf9940f7ecf2c90dc2", - "native_dump_abi": { - "event_attributes_bytes": 60, - "meta_info_bytes": 48, - "patched_hot_path": [ - "GLOBALTIMERLO", - "ENCODE_EVENT_ID", - "STORE_GLOBAL_32", - "ADD_WRITE_PTR_64_4" - ], - "placeholder": [ - "READ_CLUSTER_CTARANK", - "GLOBALTIMERLO", - "ENCODE_EVENT_ID", - "COMPUTE_SHARED_PLACEHOLDER_ADDRESS", - "STORE_SHARED_WEAK_32", - "PMEVENT" - ], - "range_attributes_bytes": 72, - "range_pop_event_id": 31, - "sentinel_event_id": 0 - }, - "patch_artifact_sha256": { - "patched_disassembly": "7a7689265c4d59fff9799f5a283e22a3ae9eb3ed91a08743106f6213931dd1c3", - "patched_kernel": "b6c57768c8bf1a1d20ad52d8987c571318340f6d2b006cb62ae58c1d934a44f4", - "unpatched_disassembly": "5d28eac4ad110174ab8fb99d1437ebf0e99b4de4df6c1cee2341e2ef210f5af7", - "unpatched_kernel": "823305d494ef1cc26146443ac9d5108da2c13684b5a07aeddc72ac04909550af" + "info": { + "legacy_iket_cubin": 0, + "magic_number": 3133075869, + "max_event_id": 31, + "max_event_name_size": 32, + "supported_features": 3, + "version_major": 0, + "version_minor": 5 + }, + "kernels": [ + { + "event_sequence": [ + { + "event_id": 1, + "event_name": "bool", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "i8", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 1, + "range_id": 0 + }, + { + "event_id": 3, + "event_name": "u8", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 2, + "range_id": 0 + }, + { + "event_id": 4, + "event_name": "i16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 3, + "range_id": 0 + }, + { + "event_id": 5, + "event_name": "u16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 4, + "range_id": 0 + }, + { + "event_id": 6, + "event_name": "i32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 5, + "range_id": 0 + }, + { + "event_id": 7, + "event_name": "u32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 6, + "range_id": 0 + }, + { + "event_id": 8, + "event_name": "i64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 7, + "range_id": 0 + }, + { + "event_id": 9, + "event_name": "u64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 16, + "range_id": 0 + }, + { + "event_id": 10, + "event_name": "f32", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 13, + "range_id": 0 + }, + { + "event_id": 11, + "event_name": "f64", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 14, + "range_id": 0 + }, + { + "event_id": 12, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 12, + "event_name": "token_payload", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 5, + "range_id": 742653351 + }, + { + "event_id": 13, + "event_name": "stack_payload", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 13, + "range_id": 3918402126 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 14, + "event_name": "filler00", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 15, + "event_name": "filler01", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 16, + "event_name": "filler02", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 17, + "event_name": "filler03", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 18, + "event_name": "filler04", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 19, + "event_name": "filler05", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 20, + "event_name": "filler06", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 21, + "event_name": "filler07", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 22, + "event_name": "filler08", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 23, + "event_name": "filler09", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 24, + "event_name": "filler10", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 25, + "event_name": "filler11", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 26, + "event_name": "filler12", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 27, + "event_name": "filler13", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 28, + "event_name": "filler14", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 29, + "event_name": "filler15", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 30, + "event_name": "filler16", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + } + ], + "kernel_name": "kernel_cutlass_native_payload_kernel_0" + } + ], + "ranges": [ + { + "color": -1, + "evt_pair_mode": 1, + "range_id": 742653351, + "range_name": "token_payload", + "range_scope": 0, + "range_type": 1 + }, + { + "color": -1, + "evt_pair_mode": 0, + "range_id": 3918402126, + "range_name": "stack_payload", + "range_scope": 0, + "range_type": 2 + } + ] + }, + "metadata_sha256": "906aae335e2fcdc13e6a600230c839a4920bec60a1b84b194eaf795f0acf94a5" + } }, "profile": { "compiler_flags": [ "iket", - "--dump-dir=", + "--dump-dir=/", "--keep-ptx", "--keep-cubin", "--keep-sass" ], "cutlass_dsl": "4.6.0", "driver": "NVIDIA B200, 595.58.03", - "instrument_method": "NativeDump", - "nvdisasm": "nvdisasm: NVIDIA (R) CUDA disassembler\nCopyright (c) 2005-2026 NVIDIA Corporation\nBuilt on Mon_Mar_02_09:52:52_PM_PST_2026\nCuda compilation tools, release 13.2, V13.2.51\nBuild cuda_13.2.r13.2/compiler.37434383_0" + "nvdisasm": "nvdisasm: NVIDIA (R) CUDA disassembler\nCopyright (c) 2005-2026 NVIDIA Corporation\nBuilt on Tue_Jun_09_02:42:28_PM_PDT_2026\nCuda compilation tools, release 13.3, V13.3.73\nBuild cuda_13.3.r13.3/compiler.38244171_0", + "nvdisasm_distribution": "13.3.73", + "nvrtc_distribution": "13.2.78" }, - "schema_version": 2, + "schema_version": 3, "wheels": { "nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl": "fa17084b07c0dca68a42892f771b4b1b40fbe9b91660209623e61cea611cae8c", "nvidia_cuda_nvrtc-13.2.78-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl": "a9049031da08cbedd0c20e3470e5a978dc330af0e0326b3b05774718c665dc3e", diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tests/python/tirx/iket/test_iket_orchestration.py index be12799f3acf..ff59170384fe 100644 --- a/tests/python/tirx/iket/test_iket_orchestration.py +++ b/tests/python/tirx/iket/test_iket_orchestration.py @@ -149,6 +149,7 @@ def test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkey assert captured["timeout"] == 12.5 assert captured["env"]["IKET_TEST_ENV"] == "present" assert captured["env"]["TVM_IKET_OFFICIAL_PROFILE"] == "cutlass-4.6.0" + assert captured["env"]["TVM_IKET_INJECTED_CHILD_ENABLE"] == "1" assert os.environ["TVM_IKET_OFFICIAL_PROFILE"] == "inherited-profile" assert result.output_dir == target assert result.command == (sys.executable, "workload.py") diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 78d766045155..0b2d2bcfc283 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -25,6 +25,7 @@ import shutil import subprocess import sys +from importlib import metadata from pathlib import Path import pytest @@ -64,18 +65,6 @@ def plain_entry(out: T.Buffer((32,), "int32")): out[tx] = tx + 7 -@T.prim_func -def explicit_cuda_shuffle_guard(out: T.Buffer((64,), "int32")): - T.device_entry() - iket = IketProfiler() - bx = T.cta_id([2]) - tx = T.thread_id([64]) - warp = T.cuda.__shfl_sync(T.uint32(0xFFFFFFFF), tx // 32, 0, 32) - if (bx == 0) & (warp == 0): - iket.mark("warp-zero") - out[tx] = tx - - @T.prim_func def push_pop_kernel(out: T.Buffer((32,), "int32")): T.device_entry() @@ -106,97 +95,124 @@ def token_loop(n: T.int32, out: T.Buffer((32,), "int32")): @T.prim_func -def overlapping_ranges(out: T.Buffer((32,), "int32")): +def payload_kernel(out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - outer = iket.range_start("overlap") - inner = iket.range_start("overlap") - iket.range_end(inner) - iket.range_end(outer) + iket.mark("payload", tx) out[tx] = tx @T.prim_func -def repeated_range_end(out: T.Buffer((32,), "int32")): +def payload_types(n: T.int64, out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - token = iket.range_start("twice") - iket.range_end(token) + iket.mark("bool", tx == 0) + iket.mark("i8", T.int8(-8)) + iket.mark("u8", T.uint8(8)) + iket.mark("i16", T.int16(-16)) + iket.mark("u16", T.uint16(16)) + iket.mark("i32", T.int32(-32)) + iket.mark("u32", T.uint32(32)) + iket.mark("i64", n) + iket.mark("u64", T.uint64(64)) + iket.mark("f32", T.float32(-3.25)) + iket.mark("f64", T.float64(6.5)) + token = iket.range_start("token_payload", T.int32(-7)) + iket.range_end(token, T.int32(9)) + iket.range_push("stack_payload", T.float32(1.5)) + iket.range_pop() + out[tx] = tx + + +@T.prim_func +def payload_presence_mismatch(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + token = iket.range_start("mismatch", tx) iket.range_end(token) out[tx] = tx @T.prim_func -def unbalanced_stack(out: T.Buffer((32,), "int32")): +def payload_type_mismatch(out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - iket.range_pop() + token = iket.range_start("mismatch", tx) + iket.range_end(token, T.uint32(tx)) out[tx] = tx @T.prim_func -def payload_kernel(out: T.Buffer((32,), "int32")): +def sentinel_only_payload(out: T.Buffer((32,), "int32")): T.device_entry() + iket = IketProfiler() tx = T.thread_id([32]) - T.evaluate(tvm.tirx.call_intrin("", "tirx.cuda.iket_mark", "payload", tx)) + token = iket.sentinel_token("not-a-declaration") + iket.range_end(token, out[tx]) out[tx] = tx @T.prim_func -def loop_carried_divergent_token(out: T.Buffer((32,), "int32")): +def payload_float16(out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - guard = T.alloc_local((1,), "int32") - guard[0] = 0 - token = iket.sentinel_token("loop") - for _i in T.serial(2, unroll=False): - token = iket.sentinel_token("loop") - if guard[0] == 0: - token = iket.range_start("loop") - iket.range_end(token) - guard[0] = tx - out[tx] = guard[0] + iket.mark("bad", T.float16(1)) + out[tx] = tx @T.prim_func -def while_carried_divergent_token(out: T.Buffer((32,), "int32")): +def payload_bfloat16(out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - guard = T.alloc_local((1,), "int32") - iteration = T.alloc_local((1,), "int32") - guard[0] = 0 - iteration[0] = 0 - token = iket.sentinel_token("while-loop") - while iteration[0] < 2: - token = iket.sentinel_token("while-loop") - if guard[0] == 0: - token = iket.range_start("while-loop") - iket.range_end(token) - guard[0] = tx - iteration[0] = iteration[0] + 1 - out[tx] = guard[0] + iket.mark("bad", T.bfloat16(1)) + out[tx] = tx + + +@T.prim_func +def payload_pointer(out: T.Buffer((32,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + T.evaluate(tvm.tirx.call_intrin("", "tirx.cuda.iket_mark", "bad", out.data)) + out[tx] = tx + + +@T.prim_func +def payload_vector(out: T.Buffer((1,), "int32x4")): + T.device_entry() + T.evaluate(tvm.tirx.call_intrin("", "tirx.cuda.iket_mark", "bad", out[0])) + + +@T.prim_func +def schema_i32(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("shared-schema", T.int32(tx)) + out[tx] = tx + + +@T.prim_func +def schema_u32(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("shared-schema", T.uint32(tx)) + out[tx] = tx @T.prim_func -def annotated_outer_loop_with_nested_break(out: T.Buffer((32,), "int32")): +def schema_no_payload(out: T.Buffer((32,), "int32")): T.device_entry() iket = IketProfiler() tx = T.thread_id([32]) - value = T.alloc_local((1,), "int32") - value[0] = 0 - for _outer in T.serial(2, unroll=False): - iket.range_push("outer") - for inner in T.serial(2, unroll=False): - if inner == 1: - break - value[0] = value[0] + 1 - iket.range_pop() - out[tx] = value[0] + iket.mark("shared-schema") + out[tx] = tx @T.prim_func @@ -297,6 +313,46 @@ def _official_global_bytes(source, symbol): return bytes(values) +def _event_bytes(source, name): + match = re.search(rf"__iket_evt_decl_{re.escape(name)}_(\d+)_attrs", source) + assert match is not None, name + event_id = int(match.group(1)) + return event_id, _official_global_bytes(source, f"__iket_evt_decl_{name}_{event_id}_attrs") + + +def _many_marks(count): + marks = "\n".join(f' iket.mark("e{index:04d}")' for index in range(count)) + source = f"""@T.prim_func +def main(out: T.Buffer((1,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([1]) +{marks} + out[tx] = 1 +""" + return tvm.script.from_source(source, {"T": T, "IketProfiler": IketProfiler}) + + +def _packaged_nvdisasm(): + distribution = metadata.distribution("nvidia-cuda-nvdisasm") + return Path(distribution.locate_file("nvidia/cu13/bin/nvdisasm")) + + +def _nvrtc_disassemble(source, tmp_path): + from tvm.support.nvcc import compile_cuda + + tmp_path.mkdir(parents=True, exist_ok=True) + cubin = compile_cuda(source, target_format="cubin", arch="sm_100a", compiler="nvrtc") + cubin_path = tmp_path / "kernel.cubin" + cubin_path.write_bytes(cubin) + return subprocess.run( + [_packaged_nvdisasm(), "-c", cubin_path], + check=True, + capture_output=True, + text=True, + ).stdout + + def test_public_interface_is_official_only(): signature = inspect.signature(IketProfiler.compile) assert "backend" not in signature.parameters @@ -311,6 +367,12 @@ def test_public_interface_is_official_only(): assert "T.tirx.iket" not in script assert tvm.script.from_source(script).script() == script + payload_script = payload_types.script() + assert 'T.cuda.iket.mark("i8", T.int8(-8))' in payload_script + assert 'T.cuda.iket.range_start("token_payload", -7)' in payload_script + assert "T.cuda.iket.range_end(token, 9)" in payload_script + assert tvm.script.from_source(payload_script).script() == payload_script + @pytest.mark.parametrize( "name", @@ -353,6 +415,34 @@ def main(out: T.Buffer((32,), "int32")): assert "iket" not in sources[1].lower() +def test_verified_injected_child_automatically_enables_plain_jit(monkeypatch): + monkeypatch.setenv("TVM_IKET_INJECTED_CHILD_ENABLE", "1") + monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.0") + monkeypatch.setenv("CUDA_INJECTION64_PATH", "/verified/libsmodel_injection.so") + monkeypatch.setenv("SMODEL_INJECTION_CONFIG", "/verified/config.json") + executable = tvm.compile(tvm.IRModule({"main": serial_a}), target=TARGET, tir_pipeline="tirx") + source = _cuda_source(executable) + assert "__iket_evt_decl_a_1_attrs" in source + + +@pytest.mark.parametrize( + "missing_env", + ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG", "TVM_IKET_INJECTED_CHILD_ENABLE"), +) +def test_injected_child_auto_enable_remains_fail_closed(monkeypatch, missing_env): + values = { + "TVM_IKET_INJECTED_CHILD_ENABLE": "1", + "TVM_IKET_OFFICIAL_PROFILE": "cutlass-4.6.0", + "CUDA_INJECTION64_PATH": "/verified/libsmodel_injection.so", + "SMODEL_INJECTION_CONFIG": "/verified/config.json", + } + for name, value in values.items(): + monkeypatch.setenv(name, value) + monkeypatch.delenv(missing_env) + executable = tvm.compile(tvm.IRModule({"main": serial_a}), target=TARGET, tir_pipeline="tirx") + assert "iket" not in _cuda_source(executable).lower() + + @pytest.mark.parametrize( ("module_name", "factory_name"), [ @@ -439,66 +529,196 @@ def test_token_sentinel_and_dynamic_alternation_lowering(): assert source.count("tvm_builtin_iket_official_event(token_ptr[0])") == 2 assert "case 1:" in source assert "case 2:" in source - assert "case 3:" in source assert "case 31:" in source - for name, event_id in (("even", 1), ("odd", 2), ("sentinel", 3)): + assert "__iket_evt_decl_sentinel" not in source + assert "__iket_range_decl_sentinel" not in source + for name, event_id in (("even", 1), ("odd", 2)): event = _official_global_bytes(source, f"__iket_evt_decl_{name}_{event_id}_attrs") assert int.from_bytes(event[4:8], "little") == event_id assert int.from_bytes(event[16:20], "little") == 4 -def test_explicit_cuda_shuffle_broadcast_is_warp_uniform(): - source = _cuda_source(_compile(explicit_cuda_shuffle_guard)) - assert "__iket_evt_decl_warp_zero_1_attrs" in source +def test_payload_metadata_and_native_record_layout(): + source = _cuda_source(_compile(payload_types)) + expected_payload_types = { + "bool": 2, + "i8": 1, + "u8": 2, + "i16": 3, + "u16": 4, + "i32": 5, + "u32": 6, + "i64": 7, + "u64": 16, + "f32": 13, + "f64": 14, + "token_payload": 5, + "stack_payload": 13, + } + for name, payload_type in expected_payload_types.items(): + _event_id, event = _event_bytes(source, name) + assert int.from_bytes(event[8:12], "little") == 3 + assert int.from_bytes(event[12:16], "little") == payload_type + + assert int.from_bytes(_event_bytes(source, "token_payload")[1][16:20], "little") == 4 + assert int.from_bytes(_event_bytes(source, "stack_payload")[1][16:20], "little") == 1 + assert "activemask.b32 %%mask" in source + assert "elect.sync _|%%p, %%mask" in source + assert "@%%p st.weak.shared.b32 [%%r+4], %%payload32" in source + assert "@%%p st.weak.shared.b64 [%%r+8], %%payload64" in source + helper_start = source.index("template ") + helper_end = source.index('extern "C" __global__', helper_start) + assert "__shfl" not in source[helper_start:helper_end] + + +def test_no_payload_native_helper_is_unchanged(): + source = _cuda_source(_compile(push_pop_kernel)) + helper = source[source.index("template ") :] + assert "activemask" not in helper + assert "elect.sync" not in helper + assert "st.weak.shared.u32 [%%r], %%t" in helper + assert "st.weak.shared.u64" not in helper -def test_nested_loop_control_does_not_reject_annotated_outer_loop(): - source = _cuda_source(_compile(annotated_outer_loop_with_nested_break)) - assert "__iket_evt_decl_outer_1_attrs" in source +def test_sentinel_only_has_no_declaration_and_guards_payload_evaluation(): + source = _cuda_source(_compile(sentinel_only_payload)) + assert "__iket_evt_decl" not in source + assert "__iket_range_decl" not in source + kernel = source[source.index("sentinel_only_payload_kernel") :] + guard = kernel.index("if (token_ptr[0] != (uint)0)") + payload_load = kernel.index("out_ptr[((int)threadIdx.x)]", guard) + event = kernel.index("tvm_builtin_iket_official_event", guard) + assert guard < event < payload_load @pytest.mark.parametrize( - "kernel_func", (loop_carried_divergent_token, while_carried_divergent_token) + ("kernel_func", "message"), + [ + pytest.param(payload_float16, "supports only", id="float16"), + pytest.param(payload_bfloat16, "supports only", id="bfloat16"), + pytest.param(payload_pointer, "scalar numeric", id="pointer"), + pytest.param(payload_vector, "scalar value", id="vector"), + ], ) -def test_unproven_warp_convergence_warns_and_compiles(kernel_func, capfd): - source = _cuda_source(_compile(kernel_func)) - warning = capfd.readouterr().err - - assert "IKET warp convergence could not be proven for event site" in warning - assert "continuing because convergence diagnostics are advisory" in warning - assert "__iket_evt_decl" in source +def test_rejects_invalid_payload_types(kernel_func, message): + with pytest.raises(TypeError, match=message): + _compile(kernel_func) @pytest.mark.parametrize( ("kernel_func", "message"), [ - pytest.param(payload_kernel, "does not support payloads", id="payload"), - pytest.param(overlapping_ranges, "strictly alternating", id="overlap"), - pytest.param(repeated_range_end, "strictly alternating", id="repeated-end"), - pytest.param(unbalanced_stack, "balanced range_push/range_pop", id="unbalanced-stack"), + pytest.param(payload_presence_mismatch, "both range_start and range_end", id="presence"), + pytest.param(payload_type_mismatch, "changes payload type", id="dtype"), ], ) -def test_rejects_unsupported_semantics(kernel_func, message): - with pytest.raises(ValueError, match=message): +def test_rejects_token_payload_schema_mismatch(kernel_func, message): + with pytest.raises((TypeError, ValueError), match=message): _compile(kernel_func) -def test_declaration_module_and_architecture_boundaries(): +def test_rejects_cross_kernel_payload_schema_conflicts(): + with pytest.raises(TypeError, match="changes payload type.*across kernels"): + IketProfiler().compile( + tvm.IRModule({"i32": schema_i32, "u32": schema_u32}), + target=TARGET, + tir_pipeline="tirx", + ) + with pytest.raises(ValueError, match="changes payload presence.*across kernels"): + IketProfiler().compile( + tvm.IRModule({"i32": schema_i32, "none": schema_no_payload}), + target=TARGET, + tir_pipeline="tirx", + ) + + +def test_native_extended_module_and_architecture_boundaries(capfd): source = _cuda_source(_compile(marks_30)) assert len(re.findall(r"__iket_evt_decl_e\d\d_\d+_attrs", source)) == 30 + assert int.from_bytes(_official_global_bytes(source, "__iket_meta_info")[12:16], "little") == 31 - with pytest.raises(ValueError, match="at most 30 declarations per kernel"): - _compile(marks_31) - with pytest.raises(ValueError, match="at most 30 distinct declarations"): + source = _cuda_source(_compile(marks_31)) + warning = capfd.readouterr().err + assert "ExtendedNativeDump" in warning + assert len(re.findall(r"__iket_evt_decl_e\d\d_\d+_attrs", source)) == 31 + assert "__iket_evt_decl_e00_64_attrs" in source + assert "__iket_evt_decl_e30_94_attrs" in source + meta = _official_global_bytes(source, "__iket_meta_info") + assert int.from_bytes(meta[12:16], "little") == 4095 + + source = _cuda_source( IketProfiler().compile( tvm.IRModule({"marks_30": marks_30, "serial_a": serial_a}), target=TARGET, tir_pipeline="tirx", ) + ) + assert ( + int.from_bytes(_official_global_bytes(source, "__iket_meta_info")[12:16], "little") == 4095 + ) with pytest.raises(ValueError, match="requires SM90 or newer"): _compile(serial_a, target=tvm.target.Target({"kind": "cuda", "arch": "sm_80"})) +def test_extended_declaration_limit(capfd): + source = _cuda_source(_compile(_many_marks(4032))) + capfd.readouterr() + assert "__iket_evt_decl_e0000_64_attrs" in source + assert "__iket_evt_decl_e4031_4095_attrs" in source + with pytest.raises(ValueError, match="at most 4032.*got 4033"): + _compile(_many_marks(4033)) + + +@pytest.mark.parametrize("arch", ("sm_90a", "sm_103a", "sm_110a", "sm_120a")) +def test_extended_payload_compile_only_architectures(arch, capfd): + executable = IketProfiler().compile( + tvm.IRModule({"marks": marks_31, "payload": payload_types}), + target=tvm.target.Target({"kind": "cuda", "arch": arch}), + tir_pipeline="tirx", + ) + capfd.readouterr() + source = _cuda_source(executable) + meta = _official_global_bytes(source, "__iket_meta_info") + assert int.from_bytes(meta[12:16], "little") == 4095 + assert "elect.sync _|%%p, %%mask" in source + + +def test_native_and_extended_payload_placeholder_sass(tmp_path): + native_no_payload_source = _cuda_source(_compile(push_pop_kernel)) + native_no_payload_sass = _nvrtc_disassemble( + native_no_payload_source, tmp_path / "native-no-payload" + ) + assert "ELECT" not in native_no_payload_sass + + native_payload_source = _cuda_source(_compile(payload_types)) + native_sass = _nvrtc_disassemble(native_payload_source, tmp_path / "native") + assert "ELECT" in native_sass + assert "STS" in native_sass + assert native_sass.count("SHFL") == native_no_payload_sass.count("SHFL") + + extended_source = _cuda_source(_compile(marks_31)) + extended_sass = _nvrtc_disassemble(extended_source, tmp_path / "extended") + assert "ELECT" not in extended_sass + assert "PMTRIG" in extended_sass + assert "STS.64" in extended_sass or "STS" in extended_sass + + extended_payload_source = _cuda_source( + IketProfiler().compile( + tvm.IRModule({"marks": marks_31, "payload": payload_types}), + target=TARGET, + tir_pipeline="tirx", + ) + ) + _event_id, event = _event_bytes(extended_payload_source, "i64") + assert int.from_bytes(event[8:12], "little") == 5 + assert "@%%p st.weak.shared.b32 [%%r+8], %%payload32" in extended_payload_source + assert "@%%p st.weak.shared.b64 [%%r+8], %%payload64" in extended_payload_source + extended_payload_sass = _nvrtc_disassemble( + extended_payload_source, tmp_path / "extended-payload" + ) + assert "ELECT" in extended_payload_sass + + @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_multi_kernel_module_has_no_tvm_control_plane(): executable = IketProfiler().compile( @@ -611,31 +831,74 @@ def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch def test_cutlass_4_6_0_oracle_manifest_integrity(): oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8")) - assert oracle["schema_version"] == 2 + assert oracle["schema_version"] == 3 assert oracle["profile"]["cutlass_dsl"] == "4.6.0" - assert oracle["profile"]["instrument_method"] == "NativeDump" - assert "--dump-dir=" in oracle["profile"]["compiler_flags"] - metadata_bytes = json.dumps(oracle["metadata"], sort_keys=True, separators=(",", ":")).encode() - assert hashlib.sha256(metadata_bytes).hexdigest() == oracle["metadata_sha256"] - abi = oracle["native_dump_abi"] + assert oracle["profile"]["nvdisasm_distribution"] == "13.3.73" + assert "V13.3.73" in oracle["profile"]["nvdisasm"] + assert "--dump-dir=/" in oracle["profile"]["compiler_flags"] + assert set(oracle["cases"]) == { + "native_no_payload", + "native_payload", + "extended_no_payload", + "extended_payload", + } + for name, case in oracle["cases"].items(): + metadata_bytes = json.dumps( + case["metadata"], sort_keys=True, separators=(",", ":") + ).encode() + assert hashlib.sha256(metadata_bytes).hexdigest() == case["metadata_sha256"] + expected_method = 3 if name.startswith("native_") else 5 + assert case["instrument_methods"] == [expected_method] + assert case["metadata"]["info"]["max_event_id"] == (31 if expected_method == 3 else 4095) + user_events = [ + event for event in case["metadata"]["events"] if event["event_id"] not in (0, 31) + ] + assert len(user_events) == (30 if expected_method == 3 else 31) + + native_payloads = { + event["event_name"]: event["payload_type"] + for event in oracle["cases"]["native_payload"]["metadata"]["events"] + } + assert ( + native_payloads + | { + "bool": 2, + "i8": 1, + "u8": 2, + "i16": 3, + "u16": 4, + "i32": 5, + "u32": 6, + "i64": 7, + "u64": 16, + "f32": 13, + "f64": 14, + } + == native_payloads + ) + + abi = oracle["abi"] assert abi["sentinel_event_id"] == 0 assert abi["range_pop_event_id"] == 31 + assert abi["native_user_event_ids"] == [1, 30] + assert abi["extended_user_event_ids"] == [64, 4095] + assert abi["max_user_declarations"] == 4032 assert ( abi["meta_info_bytes"], abi["event_attributes_bytes"], abi["range_attributes_bytes"], ) == (48, 60, 72) - assert abi["patched_hot_path"] == [ - "GLOBALTIMERLO", - "ENCODE_EVENT_ID", - "STORE_GLOBAL_32", - "ADD_WRITE_PTR_64_4", - ] + assert abi["record_layouts"]["native_payload_32"]["payload"] == [4, 4] + assert abi["record_layouts"]["native_payload_64"]["payload"] == [8, 8] + assert abi["record_layouts"]["extended_payload_32"]["payload"] == [8, 4] + assert abi["record_layouts"]["extended_payload_64"]["payload"] == [8, 8] all_hashes = [ - *oracle["artifact_sha256"].values(), - *oracle["patch_artifact_sha256"].values(), *oracle["wheels"].values(), - oracle["metadata_sha256"], + *( + digest + for case in oracle["cases"].values() + for digest in (*case["artifact_sha256"].values(), case["metadata_sha256"]) + ), ] assert all(re.fullmatch(r"[0-9a-f]{64}", value) for value in all_hashes) @@ -645,10 +908,11 @@ def test_external_trace_contract(): if trace_path is None: pytest.skip("set TVM_IKET_OFFICIAL_TRACE_JSON after the locked run-iket workload") trace = json.loads(Path(trace_path).read_text(encoding="utf-8")) - assert len(trace["launches"]) == 1 - launch = trace["launches"][0] - assert launch["kernelName"] == "canonical_iket_workload_kernel" + assert len(trace["launches"]) == 3 + launches = {launch["kernelName"]: launch for launch in trace["launches"]} strings = trace["stringTable"] + + launch = launches["canonical_iket_workload_kernel"] assert [strings[marker["markerNameIdx"]] for marker in launch["markers"]] == [ "checkpoint", "inside_stack", @@ -666,6 +930,58 @@ def test_external_trace_contract(): item["endTs"], ] + native = launches["native_payload_workload_kernel"] + markers = {strings[item["markerNameIdx"]]: item for item in native["markers"]} + assert (markers["lane_payload"]["payloadType"], markers["lane_payload"]["payloadVal"]) == ( + 5, + 100, + ) + assert ( + markers["first_active_lane"]["payloadType"], + markers["first_active_lane"]["payloadVal"], + ) == (5, 5) + assert (markers["wide_payload"]["payloadType"], markers["wide_payload"]["payloadVal"]) == ( + 7, + 0x100000000, + ) + assert ( + markers["negative_payload"]["payloadType"], + markers["negative_payload"]["payloadVal"], + ) == (5, 0xFFFFFFE0) + assert ( + markers["bool_true_payload"]["payloadType"], + markers["bool_true_payload"]["payloadVal"], + ) == (2, 1) + assert ( + markers["bool_false_payload"]["payloadType"], + markers["bool_false_payload"]["payloadVal"], + ) == (2, 0) + assert ( + markers["float32_payload"]["payloadType"], + markers["float32_payload"]["payloadVal"], + ) == (13, 0xC0500000) + assert ( + markers["float64_payload"]["payloadType"], + markers["float64_payload"]["payloadVal"], + ) == (14, 0x401A000000000000) + ranges = {strings[item["rangeNameIdx"]]: item for item in native["ranges"]} + assert [ + (event["payloadType"], event["payloadVal"]) + for event in ranges["token_payload"]["internalEvents"] + ] == [(5, 200), (5, 300)] + assert ( + ranges["stack_payload"]["internalEvents"][0]["payloadType"], + ranges["stack_payload"]["internalEvents"][0]["payloadVal"], + ) == (5, 400) + + extended = launches["extended_payload_workload_kernel"] + markers = {strings[item["markerNameIdx"]]: item for item in extended["markers"]} + assert len(markers) == 31 + assert ( + markers["extended_lane_payload"]["payloadType"], + markers["extended_lane_payload"]["payloadVal"], + ) == (5, 500) + def test_external_patch_contract(tmp_path): run_dir = os.environ.get("TVM_IKET_OFFICIAL_PATCH_RUN_DIR") @@ -682,6 +998,8 @@ def test_external_patch_contract(tmp_path): str(verifier), "--run-dir", run_dir, + "--kernel", + "canonical_iket_workload_kernel", "--nvdisasm", nvdisasm, "--output-dir", @@ -693,5 +1011,5 @@ def test_external_patch_contract(tmp_path): oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8")) assert report["schema_version"] == 1 assert report["site_count"] > 0 - assert report["normalized_signature"] == oracle["native_dump_abi"]["patched_hot_path"] + assert report["normalized_signature"] == oracle["abi"]["native_patched_hot_path"] assert all(re.fullmatch(r"[0-9a-f]{64}", value) for value in report["sha256"].values()) From 76b63de79e3ca052d65fada470a7378edaa077b5 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Tue, 4 Aug 2026 01:34:04 -0400 Subject: [PATCH 3/5] fix(lower-tirx): route IKET SASS test to GPU CI --- tests/python/tirx/iket/test_iket_profiler.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 0b2d2bcfc283..6a859e8c3d31 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -683,6 +683,8 @@ def test_extended_payload_compile_only_architectures(arch, capfd): assert "elect.sync _|%%p, %%mask" in source +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_native_and_extended_payload_placeholder_sass(tmp_path): native_no_payload_source = _cuda_source(_compile(push_pop_kernel)) native_no_payload_sass = _nvrtc_disassemble( From e2d5b681eca96d517417018c59bf2464ddf63400 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Tue, 4 Aug 2026 02:31:17 -0400 Subject: [PATCH 4/5] fix(lower-tirx): pin IKET SASS test architecture --- tests/python/tirx/iket/test_iket_profiler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 6a859e8c3d31..2e85e8b6ef54 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -685,7 +685,8 @@ def test_extended_payload_compile_only_architectures(arch, capfd): @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_native_and_extended_payload_placeholder_sass(tmp_path): +def test_native_and_extended_payload_placeholder_sass(monkeypatch, tmp_path): + monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") native_no_payload_source = _cuda_source(_compile(push_pop_kernel)) native_no_payload_sass = _nvrtc_disassemble( native_no_payload_source, tmp_path / "native-no-payload" From 5f134053d937fa7025eccf55f956a986fb959618 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Tue, 4 Aug 2026 08:24:32 -0400 Subject: [PATCH 5/5] fix(lower-tirx): skip IKET SASS test without pinned nvdisasm The GPU-marked SASS test resolves nvdisasm from the pinned nvidia-cuda-nvdisasm distribution, which ships with the CUTLASS DSL wheels and is absent from the CI GPU image. Looking it up raised PackageNotFoundError instead of skipping. Guard the test on the distribution being installed so it still runs wherever the pinned toolchain exists. Co-Authored-By: Claude Opus 5 (1M context) --- tests/python/tirx/iket/test_iket_profiler.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 2e85e8b6ef54..4ad200e912b0 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -338,6 +338,13 @@ def _packaged_nvdisasm(): return Path(distribution.locate_file("nvidia/cu13/bin/nvdisasm")) +def _has_packaged_nvdisasm(): + try: + return _packaged_nvdisasm().exists() + except metadata.PackageNotFoundError: + return False + + def _nvrtc_disassemble(source, tmp_path): from tvm.support.nvcc import compile_cuda @@ -685,6 +692,9 @@ def test_extended_payload_compile_only_architectures(arch, capfd): @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +@pytest.mark.skipif( + not _has_packaged_nvdisasm(), reason="need the pinned nvidia-cuda-nvdisasm distribution" +) def test_native_and_extended_payload_placeholder_sass(monkeypatch, tmp_path): monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1") native_no_payload_source = _cuda_source(_compile(push_pop_kernel))