chore: import upstream snapshot with attribution

python/paddle/jit/sot/profiler/kernel_stats.py

@@ -0,0 +1,236 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "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
+#
+# 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.
+
+from __future__ import annotations
+
+import atexit
+import re
+from dataclasses import dataclass
+from enum import Enum
+
+from paddle import profiler
+from paddle.base.core import tracer_event_type_to_string
+
+from ..utils.envs import ENV_ENABLE_SOT_STEP_PROFILER
+
+
+class EventVisitor:
+    def visit(self, event_node):
+        event_type_name = tracer_event_type_to_string(event_node.type)
+        visit_method_name = f"visit_{event_type_name}"
+        if not hasattr(self, visit_method_name):
+            self.generic_visit(event_node)
+            return
+        getattr(self, visit_method_name)(event_node)
+
+    def generic_visit(self, event_node):
+        for child in event_node.children_node:
+            self.visit(child)
+
+    def __call__(self, events):
+        for event_node in events.values():
+            self.visit(event_node)
+
+
+class KernelRunMode(Enum):
+    Dygraph = 1
+    Static = 2
+
+
+def safe_divide(a, b):
+    # Avoid division by zero
+    return a / b if b != 0 else 0
+
+
+@dataclass
+class KernelInfo:
+    name: str
+    run_mode: KernelRunMode
+    duration: float
+    cuda_kernels: list[CudaKernelInfo]
+
+
+@dataclass
+class CudaKernelInfo:
+    name: str
+    duration: float
+
+
+# TODO(SigureMo): Split into multiple files by visitor type
+class KernelStatsVisitor(EventVisitor):
+    SKIP_KERNEL_NAMES = {"full", "full_int_array", "shadow_feed"}
+    KERNEL_NAME_REGEX = re.compile("(?P<kernel_name>.+) kernel launch")
+
+    def __init__(self):
+        self.kernels = []
+
+    def get_kernel_name(self, event_name):
+        if match_obj := self.KERNEL_NAME_REGEX.match(event_name):
+            return match_obj.group("kernel_name")
+        raise ValueError(f"Unexpected event name: {event_name}")
+
+    def calc_kernel_count(self, mode):
+        return len(
+            [
+                kernel
+                for kernel in self.kernels
+                if (
+                    kernel.run_mode == mode
+                    and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
+                )
+            ]
+        )
+
+    def calc_kernel_duration(self, mode):
+        return sum(
+            [
+                kernel.duration
+                for kernel in self.kernels
+                if (
+                    kernel.run_mode == mode
+                    and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
+                )
+            ]
+        )
+
+    def find_all_cuda_kernels(self, host_event):
+        # TODO(SigureMo): Find a better way to find all CUDA kernels
+        return [
+            CudaKernelInfo(
+                device_event.name, device_event.end_ns - device_event.start_ns
+            )
+            for runtime_event in host_event.runtime_node
+            for device_event in runtime_event.device_node
+        ]
+
+    def visit_DygraphKernelLaunch(self, event_node):
+        duration = event_node.end_ns - event_node.start_ns
+        kernel_name = self.get_kernel_name(event_node.name)
+        all_cuda_kernels = self.find_all_cuda_kernels(event_node)
+        self.kernels.append(
+            KernelInfo(
+                kernel_name, KernelRunMode.Dygraph, duration, all_cuda_kernels
+            )
+        )
+        self.generic_visit(event_node)
+
+    def visit_StaticKernelLaunch(self, event_node):
+        duration = event_node.end_ns - event_node.start_ns
+        kernel_name = self.get_kernel_name(event_node.name)
+        all_cuda_kernels = self.find_all_cuda_kernels(event_node)
+        self.kernels.append(
+            KernelInfo(
+                kernel_name, KernelRunMode.Static, duration, all_cuda_kernels
+            )
+        )
+        self.generic_visit(event_node)
+
+    def summary(self) -> str:
+        static_kernel_duration = self.calc_kernel_duration(KernelRunMode.Static)
+        dygraph_kernel_duration = self.calc_kernel_duration(
+            KernelRunMode.Dygraph
+        )
+        static_kernel_count = self.calc_kernel_count(KernelRunMode.Static)
+        dygraph_kernel_count = self.calc_kernel_count(KernelRunMode.Dygraph)
+
+        percentage_static_kernel_count = safe_divide(
+            static_kernel_count, static_kernel_count + dygraph_kernel_count
+        )
+        percentage_static_kernel_duration = safe_divide(
+            static_kernel_duration,
+            static_kernel_duration + dygraph_kernel_duration,
+        )
+        step_summary = ""
+        step_summary += f"dygraph kernel count: {dygraph_kernel_count}\n"
+        step_summary += f"static kernel count: {static_kernel_count}\n"
+        step_summary += f"percentage static kernel count: {percentage_static_kernel_count:.2%}\n"
+
+        step_summary += f"dygraph kernel duration: {dygraph_kernel_duration / 1000000:.2f} ms\n"
+        step_summary += f"static kernel duration: {static_kernel_duration / 1000000:.2f} ms\n"
+        step_summary += f"percentage static kernel duration: {percentage_static_kernel_duration:.2%}\n"
+        return step_summary
+
+
+class SotStepProfilerGuard:
+    EXPORT_CHROME_TRACING_PATH = "./sot-chrome-tracing/"
+    STEP_CNT = 0
+    LAST_INFO_SUMMARY = None
+
+    def __init__(self, enable_kernel_stats=True, enable_chrome_tracing=False):
+        self.enable_kernel_stats = enable_kernel_stats
+        self.enable_chrome_tracing = enable_chrome_tracing
+        self.started = False
+        self.record_event = None
+        self.summary = None
+
+    def _kernel_stats(self, prof) -> str:
+        kernel_stats_visitor = KernelStatsVisitor()
+        kernel_stats_visitor(prof.profiler_result.get_data())
+        return kernel_stats_visitor.summary()
+
+    def collect_step_info_summary(self, prof) -> str:
+        summary = ""
+        if self.enable_kernel_stats:
+            summary += self._kernel_stats(prof)
+        if self.enable_chrome_tracing:
+            # If you want to export chrome tracing, you can enable this flag
+            # and view the tracing result in https://ui.perfetto.dev/#!/viewer
+            profiler.export_chrome_tracing(
+                SotStepProfilerGuard.EXPORT_CHROME_TRACING_PATH,
+                f"step_{SotStepProfilerGuard.STEP_CNT:03d}",
+            )(prof)
+        return summary
+
+    def on_trace_ready(self, prof):
+        summary = self.collect_step_info_summary(prof)
+        if SotStepProfilerGuard.STEP_CNT == 0:
+            coldstart_title = f"SOT step profiler info summary (ColdStart, step#{SotStepProfilerGuard.STEP_CNT}):"
+            coldstart_report = f"{coldstart_title}\n{summary}"
+            print(coldstart_report)
+            self.summary = coldstart_report
+        else:
+            warmup_title = f"SOT step profiler info summary (Warmup, step#{SotStepProfilerGuard.STEP_CNT}):"
+            warmup_report = f"{warmup_title}\n{summary}"
+            if SotStepProfilerGuard.LAST_INFO_SUMMARY is None:
+                atexit.register(
+                    lambda: print(SotStepProfilerGuard.LAST_INFO_SUMMARY)
+                )
+            SotStepProfilerGuard.LAST_INFO_SUMMARY = warmup_report
+            self.summary = warmup_report
+
+    def start(self):
+        if ENV_ENABLE_SOT_STEP_PROFILER.get():
+            self.profiler = profiler.Profiler(
+                targets=[
+                    profiler.ProfilerTarget.CPU,
+                    profiler.ProfilerTarget.GPU,
+                ],
+                on_trace_ready=self.on_trace_ready,
+            )
+            self.profiler.start()
+            self.started = True
+
+    def stop(self):
+        if self.started:
+            assert self.profiler is not None
+            self.profiler.stop()
+            self.profiler = None  # Avoid to hold the profiler instance
+        SotStepProfilerGuard.STEP_CNT += 1
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()