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..c1edbed2effb 100644 --- a/python/tvm/backend/cuda/iket.py +++ b/python/tvm/backend/cuda/iket.py @@ -52,24 +52,26 @@ from tvm.script import tirx as T _PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE" -_DEFAULT_PROFILE = "cutlass-4.6.1" +_INJECTED_CHILD_ENABLE_ENV = "TVM_IKET_INJECTED_CHILD_ENABLE" +_DEFAULT_PROFILE = "cutlass-4.6.0" _POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all")) _INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG") _OUTPUT_TAIL_LINES = 100 _TERMINATION_GRACE_SECONDS = 5.0 -# Hashes are SHA-256 digests of files in the public 4.6.1 wheels. Only +# Hashes are SHA-256 digests of files in the public 4.6.0 wheels. Only # ABI-independent runtime/compiler binaries are pinned so this profile works # with every Python version supported by that CUTLASS DSL release. _OFFICIAL_PROFILES = { - "cutlass-4.6.1": { + "cutlass-4.6.0": { + "nvrtc_version": (13, 2), "versions": { - "nvidia-cutlass-dsl": "4.6.1", - "nvidia-cutlass-dsl-libs-base": "4.6.1", - "nvidia-cutlass-dsl-libs-core": "4.6.1", - "nvidia-cutlass-dsl-libs-cu13": "4.6.1", + "nvidia-cutlass-dsl": "4.6.0", + "nvidia-cutlass-dsl-libs-base": "4.6.0", + "nvidia-cutlass-dsl-libs-core": "4.6.0", + "nvidia-cutlass-dsl-libs-cu13": "4.6.0", "nvidia-cuda-nvdisasm": "13.3.73", - "nvidia-cuda-nvrtc": "13.3.33", + "nvidia-cuda-nvrtc": "13.2.78", }, "files": { "nvidia-cutlass-dsl-libs-base": { @@ -92,10 +94,10 @@ }, "nvidia-cuda-nvrtc": { "nvidia/cu13/lib/libnvrtc.so.13": ( - "e51d197b3b0d2d9d850d29977423e6ac60661d429a59c440fc04e52b6fc6750a" + "c673cf3b5099d83b98a388a2bb21e5d6f481be3c4bb956e2d74c39cb714d8c63" ), - "nvidia/cu13/lib/libnvrtc-builtins.so.13.3": ( - "7394c640e5761d13d2bbcdbc4b4c5dbac7cb53cd5bc732d78f8a5cb38638e913" + "nvidia/cu13/lib/libnvrtc-builtins.so.13.2": ( + "6b1c571cc730d5fcfd57f322e1fa7e0e65de7454b2239ff6d552a09b82d47dbe" ), }, }, @@ -218,19 +220,23 @@ def _validate_run_iket_entrypoint() -> str: and item.value == "iket.cli.main:entrypoint" for item in entry_points ): - raise _profile_error("the run-iket entry point does not match CUTLASS DSL 4.6.1") + raise _profile_error("the run-iket entry point does not match the locked CUTLASS profile") return executable -def _validate_nvrtc_13_3() -> None: +def _validate_nvrtc_version(expected_version: tuple[int, int]) -> None: + expected_label = ".".join(str(part) for part in expected_version) try: from cuda.bindings import nvrtc error, major, minor = nvrtc.nvrtcVersion() except (ImportError, OSError, RuntimeError) as err: - raise _profile_error("CUDA NVRTC 13.3 is unavailable") from err - if int(error) != 0 or (int(major), int(minor)) != (13, 3): - raise _profile_error(f"CUDA NVRTC 13.3 is required, got {int(major)}.{int(minor)}") + raise _profile_error(f"CUDA NVRTC {expected_label} is unavailable") from err + actual_version = (int(major), int(minor)) + if int(error) != 0 or actual_version != expected_version: + raise _profile_error( + f"CUDA NVRTC {expected_label} is required, got {actual_version[0]}.{actual_version[1]}" + ) def _validate_official_installation(profile_name: str) -> str: @@ -265,7 +271,7 @@ def _validate_official_installation(profile_name: str) -> str: ) executable = _validate_run_iket_entrypoint() - _validate_nvrtc_13_3() + _validate_nvrtc_version(profile_config["nvrtc_version"]) return executable @@ -363,17 +369,28 @@ def __call__(self, *args, **kwargs): class IketProfiler: """TIRx annotations compiled for NVIDIA's official IKET runtime.""" - def mark(self, name: str): - T.evaluate(T.cuda.iket.mark(name)) - - def range_start(self, name: str): - return T.cuda.iket.range_start(name) - - def range_end(self, token: tvm.tirx.Expr): - T.evaluate(T.cuda.iket.range_end(token)) - - def range_push(self, name: str): - T.evaluate(T.cuda.iket.range_push(name)) + def mark(self, name: str, payload=None): + if payload is None: + T.evaluate(T.cuda.iket.mark(name)) + else: + T.evaluate(T.cuda.iket.mark(name, payload)) + + def range_start(self, name: str, payload=None): + if payload is None: + return T.cuda.iket.range_start(name) + return T.cuda.iket.range_start(name, payload) + + def range_end(self, token: tvm.tirx.Expr, payload=None): + if payload is None: + T.evaluate(T.cuda.iket.range_end(token)) + else: + T.evaluate(T.cuda.iket.range_end(token, payload)) + + def range_push(self, name: str, payload=None): + if payload is None: + T.evaluate(T.cuda.iket.range_push(name)) + else: + T.evaluate(T.cuda.iket.range_push(name, payload)) def range_pop(self): T.evaluate(T.cuda.iket.range_pop()) @@ -450,6 +467,10 @@ def _child_environment( stacklevel=3, ) child_env[_PROFILE_ENV] = profile_name + # LowerIket also requires the two run-iket injection variables before it + # honors this marker. This enables ordinary TIRx JIT compilation only in + # children started by this locked profiling entry point. + child_env[_INJECTED_CHILD_ENABLE_ENV] = "1" return child_env 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 405174435a36..b9ca652efa70 100644 --- a/src/backend/cuda/transforms/lower_iket.cc +++ b/src/backend/cuda/transforms/lower_iket.cc @@ -24,15 +24,16 @@ #include #include +#include #include #include -#include #include #include #include #include #include +#include #include #include #include @@ -48,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, @@ -62,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; @@ -156,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) { @@ -203,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}; }; @@ -214,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 " @@ -229,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; } } @@ -266,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())) { @@ -276,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"; } @@ -377,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); @@ -400,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) {} @@ -473,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)) {} @@ -952,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(); @@ -975,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) { @@ -1027,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" { @@ -1038,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); @@ -1058,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); @@ -1085,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 @@ -1109,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: @@ -1123,281 +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: - 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)) { - 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"; - } - } - 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_; - 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: @@ -1424,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); @@ -1435,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"; @@ -1457,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)) { @@ -1475,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)); @@ -1498,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); @@ -1507,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( @@ -1533,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 acc5826381f1..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,14 +6,14 @@ # "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, # 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, @@ -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 new file mode 100644 index 000000000000..ecbe88720325 --- /dev/null +++ b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_0_oracle.json @@ -0,0 +1,2414 @@ +{ + "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 + ], + "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" + }, + "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": 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 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "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 + }, + { + "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": "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 + } + ], + "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=/", + "--keep-ptx", + "--keep-cubin", + "--keep-sass" + ], + "cutlass_dsl": "4.6.0", + "driver": "NVIDIA B200, 595.58.03", + "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": 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", + "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/oracle/iket_official_cutlass_4_6_1_oracle.json b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json deleted file mode 100644 index d84d38115d73..000000000000 --- a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json +++ /dev/null @@ -1,207 +0,0 @@ -{ - "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 - } - ], - "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 - } - ], - "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" - }, - "profile": { - "compiler_flags": [ - "iket", - "--dump-dir=", - "--keep-ptx", - "--keep-cubin", - "--keep-sass" - ], - "cutlass_dsl": "4.6.1", - "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" - }, - "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" - } -} diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tests/python/tirx/iket/test_iket_orchestration.py index d72446716f4b..ff59170384fe 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,8 @@ 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 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 956f7f34bc58..4ad200e912b0 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 @@ -36,7 +37,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 @@ -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,80 +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 schema_no_payload(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("shared-schema") + out[tx] = tx @T.prim_func @@ -280,6 +313,53 @@ 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 _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 + + 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 @@ -294,6 +374,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", @@ -336,6 +422,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"), [ @@ -422,59 +536,202 @@ 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_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", "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( - 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", - ), + 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_rejects_unsupported_or_divergent_semantics(kernel_func, message): - with pytest.raises(ValueError, match=message): +def test_rejects_invalid_payload_types(kernel_func, message): + with pytest.raises(TypeError, match=message): _compile(kernel_func) -def test_declaration_module_and_architecture_boundaries(): +@pytest.mark.parametrize( + ("kernel_func", "message"), + [ + 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_token_payload_schema_mismatch(kernel_func, message): + with pytest.raises((TypeError, ValueError), match=message): + _compile(kernel_func) + + +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 + + 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 - 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( 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 + + +@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)) + 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( @@ -519,7 +776,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 +786,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,33 +842,76 @@ 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"]["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["schema_version"] == 3 + assert oracle["profile"]["cutlass_dsl"] == "4.6.0" + 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) @@ -620,10 +921,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", @@ -641,6 +943,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") @@ -657,6 +1011,8 @@ def test_external_patch_contract(tmp_path): str(verifier), "--run-dir", run_dir, + "--kernel", + "canonical_iket_workload_kernel", "--nvdisasm", nvdisasm, "--output-dir", @@ -668,5 +1024,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())