Kernel overlap verification stage, initial version

.flake8

@@ -5,5 +5,6 @@ max-complexity = 17
 
 per-file-ignores =
     src/aiu_trace_analyzer/pipeline/__init__.py:F401
+    src/aiu_trace_analyzer/verification/__init__.py:F401
     src/aiu_trace_analyzer/__init__.py:F401,F403
     src/aiu_trace_analyzer/core/acelyzer.py:C901

src/aiu_trace_analyzer/__init__.py

@@ -12,3 +12,4 @@
 from aiu_trace_analyzer.export import *
 from aiu_trace_analyzer.core import *
 from aiu_trace_analyzer.pipeline import *
+from aiu_trace_analyzer.verification import *

src/aiu_trace_analyzer/core/acelyzer.py

@@ -15,6 +15,7 @@
 import aiu_trace_analyzer.export.exporter as output
 import aiu_trace_analyzer.logger as aiulog
 import aiu_trace_analyzer.pipeline as event_pipe
+import aiu_trace_analyzer.verification as verify_pipe
 from aiu_trace_analyzer import __version__
 
 
@@ -762,12 +763,22 @@ def register_verification_functions(self,
                                         process: processor.EventProcessor,
                                         args,
                                         exporter: output.AbstractTraceExporter):
-        verification_ctx = event_pipe.VerificationContext()
-        kernel_parent_ctx = event_pipe.KernelParentVerificationContext()
+        # for a verification stage example, see verify.py
 
-        process.register_stage(callback=event_pipe.verify, context=verification_ctx)
-        process.register_stage(callback=event_pipe.kernel_parent_collect, context=kernel_parent_ctx)
+        # strict: accelerator events have no call/return relationship, so even a fully embedded
+        # event is a data problem and not a legitimate nesting
+        overlap_verification_ctx = verify_pipe.OverlapVerificationContext(strict=True)
+        kernel_parent_ctx = verify_pipe.KernelParentVerificationContext()
+
+        # register anything that verifies the raw/unsorted order of the input before this sort-stage
+        process.register_stage(callback=event_pipe.sort_events, context=event_pipe.EventSortingContext(
+            event_types=None, sortkey=self._default_sort_ts_and_rev_dur, global_sort=True))
+
+        # anything that requires a sorted event stream below this point
+        process.register_stage(callback=verify_pipe.verify_kernel_overlap, context=overlap_verification_ctx)
+        process.register_stage(callback=verify_pipe.kernel_parent_collect, context=kernel_parent_ctx)
         process.register_stage(callback=event_pipe.pipeline_barrier, context=event_pipe._main_barrier_context)
-        process.register_stage(callback=event_pipe.kernel_parent_verify, context=kernel_parent_ctx)
-        process.register_stage(callback=event_pipe.verify_cleanup)
-        process.register_stage(callback=event_pipe.verification_result_filter)
+        process.register_stage(callback=verify_pipe.kernel_parent_verify, context=kernel_parent_ctx)
+
+        # drop all regular events and only keep report-data after this stage
+        process.register_stage(callback=verify_pipe.verification_result_filter)

src/aiu_trace_analyzer/pipeline/__init__.py

@@ -111,11 +111,4 @@
 # for reference of the template, you'd do here:
 #       from aiu_trace_analyzer.pipeline.template import myprocessing
 
-from aiu_trace_analyzer.verification.verify import verify, verify_cleanup
-from aiu_trace_analyzer.verification.verify import VerificationContext
-from aiu_trace_analyzer.verification.kernel_parent_verify import (
-    KernelParentVerificationContext,
-    kernel_parent_collect,
-    kernel_parent_verify
-)
-from aiu_trace_analyzer.verification.report import verification_result_filter
+# NOTE: the verification stages live in aiu_trace_analyzer.verification

src/aiu_trace_analyzer/pipeline/overlap.py

@@ -3,7 +3,10 @@
 import copy
 
 import aiu_trace_analyzer.logger as aiulog
-from aiu_trace_analyzer.pipeline import AbstractContext, AbstractHashQueueContext, TwoPhaseWithBarrierContext
+from aiu_trace_analyzer.pipeline import (
+    AbstractContext,
+    AbstractHashQueueContext,
+    TwoPhaseWithBarrierContext)
 from aiu_trace_analyzer.types import TraceEvent, GlobalIngestData, TraceWarning
 from aiu_trace_analyzer.pipeline.tools import PipelineContextTool
 
@@ -32,17 +35,27 @@ class OverlapDetectionContext(TwoPhaseWithBarrierContext):
        is the time stamp indicating until what time it's active
      * need to keep a list of active end ts because conflicts might happen
        within non-critical nested overlapping events
+     * strict mode: any event that starts while another event of the same stream is still
+       running is an overlap. The default (non-strict) mode accepts fully embedded events
+       as a legitimate nesting (e.g. a parent/child function-call relationship) and only
+       flags partial overlaps.
     '''
     OVERLAP_RESOLVE_DROP = 1
     OVERLAP_RESOLVE_TID = 2
     OVERLAP_RESOLVE_ASYNC = 3
     OVERLAP_RESOLVE_WARN = 4
     OVERLAP_RESOLVE_SHIFT = 5
 
+    # default event fields that define the identity of a stream/queue. Derived classes with a static
+    # partitioning can override this by assignment (never in-place: the list is shared by all
+    # instances). Hierarchical keys are supported, e.g. "args.stream".
+    _QUEUE_ID_KEYS = ["pid", "tid"]
+
     def __init__(self,
                  overlap_resolve=OVERLAP_RESOLVE_DROP,
                  ts_shift_threshold=0.0,
                  max_tid_streams=5,
+                 strict=False,
                  ) -> None:
         super().__init__(warnings=[
             TraceWarning(
@@ -57,14 +70,25 @@ def __init__(self,
         self.ts_shift_threshold = ts_shift_threshold
         self.tid_space = {}
         self.max_tid_streams = max_tid_streams
+        self.strict = strict
+        self.queue_id_keys = None  # resolved from the first event, see _select_queue_id_keys()
+
+    # select the event fields that define a stream/queue identity. Called only for the first event,
+    # so the selection may depend on the input dialect (assumed constant for one program execution).
+    # Derived classes can override, e.g. `return super()._select_queue_id_keys(event) + ["args.stream"]`
+    def _select_queue_id_keys(self, _event: TraceEvent) -> list[str]:
+        return self._QUEUE_ID_KEYS
 
     # search for events within the same pid/tid
     # accumulate a queue of events for each pid/tid
     # once the queue is full, run detection and emit events that are fine
     def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]:
 
         tid = event["tid"] if "tid" in event else 0
-        queue_id = self.event_data_hash(event, ["pid", "tid"], ignore_missing=True)
+        if self.queue_id_keys is None:
+            self.queue_id_keys = self._select_queue_id_keys(event)
+            aiulog.log(aiulog.DEBUG, "POD: stream identity keys:", self.queue_id_keys)
+        queue_id = self.event_data_hash(event, self.queue_id_keys, ignore_missing=True)
         if queue_id not in self.queues:
             self.queues[queue_id] = (0.0, False, [])
 
@@ -78,13 +102,23 @@ def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]:
         event_ts = event["ts"]
         event_end = round(event["ts"] + event["dur"], 4)
 
+        # retire the events that have ended before this one starts, so that the blocked status
+        # refers to the ts of this event (every event leaves its own end-ts behind, so a status
+        # that's only updated after the check would keep any queue blocked forever)
+        # in SHIFT-mode, recursion happens within the same queue_id -> the update has to happen later
+        if self.overlap_resolve != self.OVERLAP_RESOLVE_SHIFT:
+            self.update_queue_status(event_ts, queue_id)
+            _, blocked, end_ts = self.queues[queue_id]
+
         aiulog.log(aiulog.TRACE, "POD queue before: ", queue_id, "from", event["pid"], tid, self.queues[queue_id])
         assert (blocked and len(end_ts) > 0) or (not blocked and len(end_ts) == 0)
         if not blocked:
             self.queues[queue_id] = (event_ts, True, [event_end])
             revents = [event]
         else:
-            if self.check_overlap_condition(event_ts, event_end, self.queues[queue_id]):
+            # a blocked queue is all it takes in strict mode: embedded events are overlaps too,
+            # so there's no additional condition to check
+            if self.strict or self.check_overlap_condition(event_ts, event_end, self.queues[queue_id]):
                 # actual overlap
                 revents = self.handle_overlap(event, queue_id)
             else:
@@ -95,7 +129,8 @@ def overlap_detection(self, event: TraceEvent) -> list[TraceEvent]:
         if self.overlap_resolve == self.OVERLAP_RESOLVE_ASYNC:
             aevents = self.update_async_event_queue(queue_id, None, event_ts)
             revents = aevents + revents  # prepend any async events that need to be injected
-        self.update_queue_status(event_ts, queue_id)
+        elif self.overlap_resolve == self.OVERLAP_RESOLVE_SHIFT:
+            self.update_queue_status(event_ts, queue_id)  # shift-mode requires update _after_ resolution
         aiulog.log(aiulog.TRACE, "POD queue after: ", queue_id, "from", event["pid"], tid, self.queues[queue_id])
         return revents
 
@@ -110,10 +145,10 @@ def check_overlap_condition(self, ts, end, qstate: OverlapTracking) -> bool:
             aiulog.log(aiulog.TRACE, "POD overlap detected", qstate, ts, end)
         return overlap
 
-    # remove only keep entries of end timestamps that are later than the new current head
+    # remove only entries of end timestamps that are strictly later than the new current head
     def update_queue_status(self, new_current: float, queue_id: int):
         end_q = self.queues[queue_id][2]
-        new_end_q = list(filter(lambda x: x >= new_current, end_q))
+        new_end_q = list(filter(lambda x: x > new_current, end_q))
         is_blocked = (len(new_end_q) > 0)  # unblock the queue if no more end-ts are remaining
         self.queues[queue_id] = (new_current, is_blocked, new_end_q)
 
@@ -138,17 +173,22 @@ def find_next_tid(self, event: TraceEvent) -> int:
         new_tid = self.tid_space[event["pid"]][event["tid"]]
         return new_tid
 
+    # single funnel for recording a detected overlap: derived classes can override this to
+    # attach per-event detail to the warning (the offending event is not available in issue_warning())
+    def _record_overlap(self, _oevent: TraceEvent) -> None:
+        self.issue_warning("overlaps")
+
     # solve a detected overlap between a pair of pairs
     def handle_overlap(self,
                        oevent: TraceEvent,
                        queue_id: int) -> list[TraceEvent]:
         if self.overlap_resolve == self.OVERLAP_RESOLVE_DROP:
             aiulog.log(aiulog.WARN, "Solving overlap conflict by dropping:", oevent)
-            self.issue_warning("overlaps")
+            self._record_overlap(oevent)
             return []
         elif self.overlap_resolve == self.OVERLAP_RESOLVE_WARN:
             aiulog.log(aiulog.WARN, "Detected overlap conflict: ", oevent["name"])
-            self.issue_warning("overlaps")
+            self._record_overlap(oevent)
             return [oevent]
         elif self.overlap_resolve == self.OVERLAP_RESOLVE_SHIFT:
             ts_shift = self.get_overlap_time(oevent["ts"], oevent["ts"]+oevent["dur"], self.queues[queue_id])
@@ -173,21 +213,22 @@ def handle_overlap(self,
                            "us: increase threshold or use different overlap res option.")
                 rlist = [oevent]
 
-            self.issue_warning("overlaps")
+            self._record_overlap(oevent)
             return rlist
         elif self.overlap_resolve == self.OVERLAP_RESOLVE_TID:
             oevent["tid"] = self.find_next_tid(oevent)
             # feed offending event back into the detector with the new TID to make sure
             # there are no collisions there either
             rlist = self.overlap_detection(oevent)
-            self.issue_warning("overlaps")
+            self._record_overlap(oevent)
             return rlist
         elif self.overlap_resolve == self.OVERLAP_RESOLVE_ASYNC:
+            # record before the event is converted to an async b/e pair (which drops its 'dur')
+            self._record_overlap(oevent)
             oevent["id"] = self.async_id
             end_ts = oevent["ts"] + oevent["dur"]
             oevent.pop("dur")
             self.async_id += 1
-            self.issue_warning("overlaps")
 
             e_event = copy.deepcopy(oevent)
             oevent["ph"] = "b"

src/aiu_trace_analyzer/profiles/verification.json

@@ -1,9 +1,10 @@
 {
     "stages": [
+        {"sort_events": true},
+        {"verify_kernel_overlap": true},
         {"kernel_parent_collect": true},
         {"pipeline_barrier": true},
         {"kernel_parent_verify": true},
-        {"verify_cleanup": true},
         {"verification_result_filter": true}
     ]
 }

src/aiu_trace_analyzer/verification/__init__.py

@@ -0,0 +1,28 @@
+# Copyright 2024-2026 IBM Corporation
+
+'''
+Verification pipeline stages and contexts.
+
+The stages of the verification pipeline are collected here, separate from the regular processing
+pipeline of aiu_trace_analyzer.pipeline. The dependency between the two is one-directional:
+verification stages build on the pipeline infrastructure (contexts, detection algorithms), the
+regular pipeline never imports anything from here.
+
+A verification stage consists of a context derived from AbstractVerificationContext (which emits
+the accumulated findings as verification meta-data events at drain time) plus one or more callback
+functions. For a minimal example of both, see verify.py.
+'''
+
+# import the verification contexts:
+from aiu_trace_analyzer.verification.verify import VerificationContext
+from aiu_trace_analyzer.verification.kernel_parent_verify import KernelParentVerificationContext
+from aiu_trace_analyzer.verification.overlap_verify import OverlapVerificationContext
+
+# import the verification stage callbacks:
+from aiu_trace_analyzer.verification.verify import verify, verify_cleanup
+from aiu_trace_analyzer.verification.kernel_parent_verify import (
+    kernel_parent_collect,
+    kernel_parent_verify
+)
+from aiu_trace_analyzer.verification.overlap_verify import verify_kernel_overlap
+from aiu_trace_analyzer.verification.report import verification_result_filter

src/aiu_trace_analyzer/verification/kernel_parent_verify.py

@@ -44,8 +44,8 @@
 
 import aiu_trace_analyzer.logger as aiulog
 from aiu_trace_analyzer.types import TraceEvent, TraceWarning
-from aiu_trace_analyzer.pipeline import TwoPhaseWithBarrierContext, AbstractContext
-from aiu_trace_analyzer.pipeline.context import AbstractVerificationContext
+from aiu_trace_analyzer.pipeline.barrier import TwoPhaseWithBarrierContext
+from aiu_trace_analyzer.pipeline.context import AbstractContext, AbstractVerificationContext
 
 
 class KernelParentVerificationContext(AbstractVerificationContext, TwoPhaseWithBarrierContext):

src/aiu_trace_analyzer/verification/overlap_verify.py

@@ -0,0 +1,101 @@
+# Copyright 2024-2026 IBM Corporation
+
+"""
+Overlap verification for accelerator compute events.
+
+Compute events of the same stream cannot overlap in time: the stream processes them one after
+another. Any detected overlap therefore indicates a data problem (e.g. inaccurate timestamps)
+and is reported as an error-level finding of the verification report.
+
+This module reuses the overlap detection of the regular processing pipeline
+(OverlapDetectionContext) and only replaces the parts that differ in verification mode:
+  * overlaps are reported, never resolved
+  * the finding is error-level and records each offending event as an instance
+  * the streams are identified per input dialect instead of by pid+tid
+"""
+
+from aiu_trace_analyzer.types import TraceEvent, TraceWarning
+from aiu_trace_analyzer.pipeline.context import AbstractContext, AbstractVerificationContext
+from aiu_trace_analyzer.pipeline.overlap import OverlapDetectionContext
+from aiu_trace_analyzer.pipeline.tools import PipelineContextTool
+
+
+class OverlapVerificationContext(OverlapDetectionContext, AbstractVerificationContext):
+    '''
+    Verification-mode variant of the overlap detection: it only reports overlapping compute
+    events, it never resolves them (always OVERLAP_RESOLVE_WARN). Every detected overlap is
+    recorded as an instance of an error-level finding and the accumulated findings are emitted
+    as verification meta-data events by AbstractVerificationContext.drain() (reached via the MRO,
+    the parent drain chains up to it; no two-phase barrier required for this context).
+
+    Compute streams are identified per dialect: FLEX events have a single stream per pid, while
+    TORCH events separate the streams within a pid by their 'args.stream' entry.
+    '''
+    test_name = "Compute Overlap Check"
+
+    _OVERLAP_WARNING = "overlaps"
+    _STREAM_KEY = "stream"
+    _STREAM_ARG = "args." + _STREAM_KEY
+    _DEFAULT_STREAM = 0  # convention: actual stream numbers start at 1, 0 means 'no stream entry'
+
+    def __init__(self, strict=False) -> None:
+        super().__init__(self.OVERLAP_RESOLVE_WARN, strict=strict)
+        # replace the resolution-oriented warning of the parent: in verification mode nothing is
+        # resolved, a detected overlap is a finding that has to fail the test
+        self.add_warning(
+            TraceWarning(
+                name=self._OVERLAP_WARNING,
+                text="Overlapping accelerator events detected: {d[count]}",
+                data={"count": 0},
+                is_error=True,
+            )
+        )
+
+    def _select_queue_id_keys(self, event: TraceEvent) -> list[str]:
+        dialect = PipelineContextTool.get_dialect_of_event(event)
+        assert dialect is not None, \
+            "OVL: cannot determine the dialect of the first event." \
+            " Register this stage before any stage that removes the jobhash."
+        if dialect.get("NAME") == "TORCH":
+            return ["pid", self._STREAM_ARG]
+        return ["pid"]
+
+    def _record_overlap(self, oevent: TraceEvent) -> None:
+        super()._record_overlap(oevent)
+        self.warnings[self._OVERLAP_WARNING].add_instance({
+            "name": oevent["name"],
+            "pid": oevent["pid"],
+            "tid": oevent["tid"],
+            "stream": oevent["args"].get(self._STREAM_KEY),
+            "ts": oevent["ts"],
+            "dur": oevent["dur"],
+        })
+
+    def add_default_stream(self, event: TraceEvent) -> bool:
+        '''
+        the stream-based separation of queues requires the stream entry to exist. Events without
+        one are all attributed to the same default stream. Returns whether a default was added.
+        '''
+        if self._STREAM_KEY in event["args"]:
+            return False
+        event["args"][self._STREAM_KEY] = self._DEFAULT_STREAM
+        return True
+
+    def remove_default_stream(self, event: TraceEvent) -> None:
+        del event["args"][self._STREAM_KEY]
+
+
+def verify_kernel_overlap(event: TraceEvent, context: AbstractContext) -> list[TraceEvent]:
+    assert isinstance(context, OverlapVerificationContext)
+
+    # only compute events are checked; anything else just passes through
+    if event["ph"] not in "X" or not PipelineContextTool.is_acc_event(event):
+        return [event]
+
+    # the default is only required to determine the queue/stream of this event: drop it again to
+    # keep the event unchanged for any downstream stage
+    default_added = context.add_default_stream(event)
+    revents = context.overlap_detection(event)
+    if default_added:
+        context.remove_default_stream(event)
+    return revents