EventReplay: extend beyond aten ops with auto-import and custom initializers

TraceLens/EventReplay/__init__.py

@@ -0,0 +1,23 @@
+###############################################################################
+# Copyright (c) 2024 - 2026 Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+###############################################################################
+
+from .event_replay import EventReplayer
+from .custom_inits import (
+    CustomInit,
+    PagedAttentionInit,
+    MoeRoutingInit,
+    extract_batch_context,
+)
+from .utils import benchmark_func
+
+__all__ = [
+    "EventReplayer",
+    "CustomInit",
+    "PagedAttentionInit",
+    "MoeRoutingInit",
+    "extract_batch_context",
+    "benchmark_func",
+]

TraceLens/EventReplay/batched_replay.py

@@ -9,7 +9,42 @@
 import argparse
 import sys
 import torch
-from utils import TensorCfg, build_tensor, benchmark_func
+
+try:
+    from .utils import TensorCfg, build_tensor, benchmark_func
+    from .event_replay import _resolve_op_func as _resolve_op_func_sdk
+except ImportError:
+    from utils import TensorCfg, build_tensor, benchmark_func
+
+    try:
+        from event_replay import _resolve_op_func as _resolve_op_func_sdk
+    except ImportError:
+        _resolve_op_func_sdk = None
+
+
+def resolve_replay_func(op_name: str):
+    """Return a callable for *op_name*, or raise RuntimeError.
+
+    Prefers EventReplayer's resolver (JIT → torch.ops → module → auto-import).
+    If this script is used as a standalone zip without ``event_replay.py``,
+    falls back to JIT and treats a ``None`` return as a miss (do not call it).
+    """
+    if _resolve_op_func_sdk is not None:
+        func, _source, _resolved = _resolve_op_func_sdk(op_name)
+        if func is None or not callable(func):
+            raise RuntimeError(f"Cannot resolve op '{op_name}'")
+        return func
+
+    try:
+        func, _ = torch._C._jit_get_operation(op_name)
+    except Exception as e:
+        raise RuntimeError(f"Cannot resolve op '{op_name}': {e}") from e
+    if func is None or not callable(func):
+        raise RuntimeError(
+            f"Cannot resolve op '{op_name}' (JIT returned {func!r}). "
+            "Place event_replay.py next to this script for custom-op resolution."
+        )
+    return func
 
 
 def _get_args_kwargs_from_ir(
@@ -99,18 +134,25 @@ def _get_args_kwargs_from_ir(
     replayed_count = 0
     errors = 0
 
-    for i, repro_info in enumerate(repro_data_list):
+    ops_to_replay = repro_data_list
+    if args.op_filter:
+        ops_to_replay = [r for r in ops_to_replay if args.op_filter in r["op_name"]]
+    if args.op_limit:
+        ops_to_replay = ops_to_replay[: args.op_limit]
+
+    for i, repro_info in enumerate(ops_to_replay):
 
         op_name = repro_info["op_name"]
         replay_ir = repro_info["replay_ir"]
-        print(f"\n[{replayed_count + 1}/{len(repro_data_list)}] Replaying: {op_name}")
+        print(f"\n[{replayed_count + 1}/{len(ops_to_replay)}] Replaying: {op_name}")
 
         # Get the PyTorch operation function
         try:
-            func, _ = torch._C._jit_get_operation(op_name)
+            func = resolve_replay_func(op_name)
         except Exception as e:
             print(
-                f"  Error: Could not find PyTorch operation '{op_name}'. Is the PyTorch version compatible? Error: {e}"
+                f"  Error: Could not find PyTorch operation '{op_name}'. "
+                f"Is the op library imported? Error: {e}"
             )
             if args.stop_on_error:
                 raise
@@ -151,15 +193,18 @@ def _get_args_kwargs_from_ir(
             errors += 1
             continue
         # --- Benchmark the function ---
-        mean_time_us = benchmark_func(
+        metrics = benchmark_func(
             lambda: func(*pos_args, **kwargs), args.device, warmup=50, avg_steps=100
         )
-        print(f"  Average time taken: {mean_time_us:.2f} microseconds")
+        mean_time_us = metrics["mean_us"]
+        print(
+            f"  Average time taken: {mean_time_us:.2f} us  (median: {metrics['median_us']:.2f} us)"
+        )
         if "count" in repro_info:
             count_workload = repro_info["count"]
             total_time_us = mean_time_us * count_workload
             print(f"  Count in workload: {count_workload}")
-            print(f"  Est time in workload: {total_time_us:.2f} microseconds")
+            print(f"  Est time in workload: {total_time_us:.2f} us")
         # --- Optionally sync again ---
         if args.device == "cuda":
             torch.cuda.synchronize()
@@ -190,7 +235,9 @@ def _get_args_kwargs_from_ir(
     print("\n--- Replay Summary ---")
     print(f"Total operations in file: {len(repro_data_list)}")
     if args.op_filter:
-        print(f"Filter applied: '{args.op_filter}'")
+        print(f"Filter applied: '{args.op_filter}' ({len(ops_to_replay)} matched)")
+    if args.op_limit:
+        print(f"Limit applied: {args.op_limit}")
     print(f"Attempted replays: {replayed_count}")
     print(f"Successful replays: {replayed_count - errors}")
     print(f"Errors encountered: {errors}")