Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / api.py: 92%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Generic MegaKernel integration without imports of concrete diagnostic tools.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable, Sequence 

20from typing import Any 

21 

22import torch 

23 

24from hyper_multicore.diagnostics.common.call import DiagnosticCall, LaunchContext 

25from hyper_multicore.diagnostics.common.graph import DiagnosticSpec, GraphDiagnosticMetadata 

26from hyper_multicore.diagnostics.common.runtime import PreparedDiagnosticRuntime 

27from hyper_multicore.diagnostics.common.session import _get_active_session 

28from hyper_multicore.scheduler.config import RuntimeConfigC, TaskDescC 

29from hyper_multicore.scheduler.graph import ComputeGraph, OperatorNode 

30 

31 

32class MegaKernelDiagnostics: 

33 """A reusable normal runtime and optional tool-independent launch boundary.""" 

34 

35 def __init__(self, runtime: PreparedDiagnosticRuntime) -> None: 

36 """Bind a normal runtime without evaluating diagnostic graph semantics.""" 

37 self._runtime = runtime 

38 

39 @property 

40 def runtime_config(self) -> torch.Tensor: 

41 """Return the original normal execution configuration.""" 

42 return self._runtime.normal_tensor 

43 

44 @property 

45 def is_active(self) -> bool: 

46 """Return whether an explicit diagnostic session is active.""" 

47 session = _get_active_session() 

48 return session is not None and session.is_capture_enabled() 

49 

50 @property 

51 def metadata(self) -> GraphDiagnosticMetadata: 

52 """Resolve shared graph metadata explicitly for diagnostics or inspection.""" 

53 return self._runtime.metadata 

54 

55 def launch( 

56 self, parameters: Sequence[torch.Tensor | None], *, direction: str, 

57 fallback_event_counters: torch.Tensor | None = None, 

58 ) -> DiagnosticCall: 

59 """Select exactly one tool, or return an ordinary no-op context. 

60 

61 Kernel adapters should bypass this method when is_active is false so 

62 ordinary execution need not create even a Host-side diagnostic call. 

63 

64 Args: 

65 parameters: Complete native ABI bindings for this invocation. 

66 direction: Forward or backward invocation direction. 

67 fallback_event_counters: Existing execution-owned event counters, if any. 

68 """ 

69 context = LaunchContext(self._runtime, parameters, direction, fallback_event_counters) 

70 session = _get_active_session() 

71 if session is None or not session.is_capture_enabled(): 

72 return DiagnosticCall(context.ordinary_arguments()) 

73 if len(parameters) != self._runtime.parameter_count: 

74 raise ValueError(f"launch expects {self._runtime.parameter_count} ABI parameters, got {len(parameters)}") 

75 if direction not in ("forward", "backward"): 

76 raise ValueError(f"invalid MegaKernel direction {direction!r}") 

77 return session.prepare_call(context) 

78 

79 

80def prepare_mega_kernel_diagnostics( 

81 graph: ComputeGraph | None, runtime_config: RuntimeConfigC, *, kernel_name: str, 

82 parameter_count: int, device: torch.device | str, rank: int = 0, owner_label: str = "Owner", 

83 owner_resolver: Callable[[OperatorNode, TaskDescC, Any], int | None] | None = None, 

84 owner_context: Any = None, 

85 execution_mode: str = "event_driven", vector_worker_stride: int = 2, 

86 graph_adapter: Callable[[Any, RuntimeConfigC], ComputeGraph] | None = None, 

87) -> MegaKernelDiagnostics: 

88 """Prepare only normal execution storage and defer all diagnostic graph work. 

89 

90 Args: 

91 graph: Completed scheduling graph; may be None if graph_adapter builds 

92 the graph of a native broadcast pipeline on first capture. 

93 runtime_config: Filled graph-sized Host configuration. 

94 kernel_name: Stable name shared by both tools. 

95 parameter_count: Number of launch binding entries. Device task positions 

96 must match input_list exactly; zero-task Host boundary nodes may 

97 describe results produced outside the MegaKernel. 

98 device: Device receiving the ordinary runtime configuration. 

99 rank: Distributed rank recorded by diagnostics. 

100 owner_label: Business display label, such as Expert, Head or Block. 

101 owner_resolver: Optional task-owner callback, never called on the normal path. 

102 owner_context: Business context passed unchanged to the owner callback. 

103 execution_mode: Event-driven queues or a broadcast ``aiv_pipeline``. 

104 vector_worker_stride: Two for paired MoE workers, one for all AIV workers. 

105 graph_adapter: Optional lazy business-owned binding of graph ports to the 

106 real native ABI. It receives a private Host config; ordinary storage 

107 and the caller's graph must remain unchanged. 

108 

109 Returns: 

110 Adapter used by any MegaKernel; dynamic view resolvers remain dump-owned. 

111 """ 

112 if graph is None and graph_adapter is None: 

113 raise ValueError("a compute graph or graph_adapter is required") 

114 if graph_adapter is not None and not callable(graph_adapter): 

115 raise TypeError("graph_adapter must be callable or None") 

116 resolved_device = torch.device(device) 

117 device_id = resolved_device.index 

118 if device_id is None: 

119 device_id = 0 if resolved_device.type == "cpu" else torch.get_device_module(resolved_device).current_device() 

120 resolved_device = torch.device(resolved_device.type, device_id) 

121 

122 def tensor_factory(data: bytes) -> torch.Tensor: 

123 """Copy one runtime image onto the execution device. 

124 

125 Args: 

126 data: Serialized runtime configuration bytes. 

127 """ 

128 return torch.frombuffer(bytearray(data), dtype=torch.uint8).clone().to(device=resolved_device) 

129 

130 if execution_mode not in ("event_driven", "aiv_pipeline"): 

131 raise ValueError(f"invalid execution mode: {execution_mode!r}") 

132 if vector_worker_stride not in (1, 2): 

133 raise ValueError("vector_worker_stride must be 1 or 2") 

134 runtime = PreparedDiagnosticRuntime( 

135 graph, runtime_config, DiagnosticSpec( 

136 kernel_name, owner_label, owner_resolver, owner_context, 

137 execution_mode, vector_worker_stride, graph_adapter, 

138 ), 

139 parameter_count=parameter_count, tensor_factory=tensor_factory, rank=rank, device_id=device_id, 

140 ) 

141 return MegaKernelDiagnostics(runtime)