Add MoE route-trace replay and timing validation

justfile

@@ -132,6 +132,9 @@ aten-emulate nickname *args:
 test-sliced-layer-builder:
     python3 transactional_emulator/testbench/test_sliced_layer_builder.py
 
+# Rust emulator timing smoke: unit tests + deterministic simulation-cycle goldens.
+timing-smoke:
+    bash transactional_emulator/testbench/timing_goldens/run_timing_smoke.sh
 
 # Unit tests for LUI+ADDI large immediate fix in ASM templates
 test-large-immediate:
@@ -171,4 +174,3 @@ multilayer-decoder-profile model="smolvlm2":
 # ATen-backed sliced emulator check: PlenaCompiler + ops.* -> emulator -> numerical check
 test-sliced-aten-emulator model="AICrossSim/clm-60m" seq_len="64" num_layers="1":
     cd PLENA_Compiler && PYTHONPATH=".:../PLENA_Tools:../transactional_emulator/testbench:..:" python3 -m compiler.aten.sliced_emulator_runner {{model}} --seq-len {{seq_len}} --num-layers {{num_layers}}
-

transactional_emulator/src/accelerator/dispatch.rs

@@ -7,11 +7,12 @@ use half::bf16;
 use quantize::MxDataType;
 
 use crate::runtime_config::{
-    HLEN, MATRIX_KV_TYPE, MATRIX_WEIGHT_TYPE, MLEN, PREFETCH_M_AMOUNT, PREFETCH_V_AMOUNT,
+    BLEN, HLEN, MATRIX_KV_TYPE, MATRIX_WEIGHT_TYPE, MLEN, PREFETCH_M_AMOUNT, PREFETCH_V_AMOUNT,
     SCALAR_FP_BASIC_CYCLES, SCALAR_FP_EXP_CYCLES, SCALAR_FP_RECI_CYCLES, SCALAR_FP_SQRT_CYCLES,
     SCALAR_INT_BASIC_CYCLES, STORE_V_AMOUNT, VECTOR_ACTIVATION_TYPE, VECTOR_KV_TYPE, VLEN,
 };
 use crate::stage_profile::{ResourceKind, StageProfiler};
+use crate::timing_overlay::{SramRange, SramSpace, TimingAccess, TimingOverlay};
 use crate::{cycle, dma, op};
 use runtime::Executor;
 
@@ -54,13 +55,20 @@ impl Accelerator {
         &mut self,
         ops: &[op::Opcode],
         mut stage_profiler: Option<&mut StageProfiler>,
+        mut timing_overlay: Option<&mut TimingOverlay>,
     ) {
         let mut pc: usize = 0; // Program counter
 
         while pc < ops.len() {
             let executed_pc = pc;
             let op = &ops[pc];
-            let profile_start_instant = if stage_profiler.is_some() {
+            let capture_elapsed = stage_profiler.is_some() || timing_overlay.is_some();
+            let timing_access = if timing_overlay.is_some() {
+                Some(self.timing_access_for_opcode(op))
+            } else {
+                None
+            };
+            let profile_start_instant = if capture_elapsed {
                 Some(Executor::current().now())
             } else {
                 None
@@ -625,35 +633,118 @@ impl Accelerator {
                 pc += 1;
             }
 
-            if let (Some(start_instant), Some(profiler)) =
-                (profile_start_instant, stage_profiler.as_deref_mut())
-            {
+            if let Some(start_instant) = profile_start_instant {
                 let elapsed_duration = Executor::current().now() - start_instant;
                 let elapsed_secs = elapsed_duration.as_picos() as f64 / 1_000_000_000_000.0;
                 let elapsed_cycles = StageProfiler::duration_to_cycles(elapsed_duration);
-                let (hbm_bytes_read, hbm_bytes_written) = if let (Some(before), Some(after)) =
-                    (profile_start_hbm, self.hbm.statistics())
+                if let (Some(access), Some(overlay)) =
+                    (timing_access, timing_overlay.as_deref_mut())
                 {
-                    (
-                        after
-                            .total_bytes_read
-                            .saturating_sub(before.total_bytes_read),
-                        after
-                            .total_bytes_written
-                            .saturating_sub(before.total_bytes_written),
-                    )
-                } else {
-                    (0, 0)
-                };
-                profiler.record(
-                    executed_pc,
-                    elapsed_secs,
-                    elapsed_cycles,
-                    resource_kind_for_opcode(op),
-                    hbm_bytes_read,
-                    hbm_bytes_written,
-                );
+                    overlay.record(access, elapsed_cycles);
+                }
+                if let Some(profiler) = stage_profiler.as_deref_mut() {
+                    let (hbm_bytes_read, hbm_bytes_written) = if let (Some(before), Some(after)) =
+                        (profile_start_hbm, self.hbm.statistics())
+                    {
+                        (
+                            after
+                                .total_bytes_read
+                                .saturating_sub(before.total_bytes_read),
+                            after
+                                .total_bytes_written
+                                .saturating_sub(before.total_bytes_written),
+                        )
+                    } else {
+                        (0, 0)
+                    };
+                    profiler.record(
+                        executed_pc,
+                        elapsed_secs,
+                        elapsed_cycles,
+                        resource_kind_for_opcode(op),
+                        hbm_bytes_read,
+                        hbm_bytes_written,
+                    );
+                }
+            }
+        }
+    }
+
+    fn timing_access_for_opcode(&self, op: &op::Opcode) -> TimingAccess {
+        let matrix_tile = *MLEN * *MLEN;
+        let vector_tile = *VLEN;
+        let matrix_prefetch = *MLEN * *PREFETCH_M_AMOUNT;
+        let vector_prefetch = *VLEN * *PREFETCH_V_AMOUNT;
+        let matrix_vector_batch = *MLEN * *BLEN;
+
+        let matrix = |start: u32, len: u32| SramRange::new(SramSpace::Matrix, start, len);
+        let vector = |start: u32, len: u32| SramRange::new(SramSpace::Vector, start, len);
+        let gp = |reg: u8| self.reg_file.read_gp(reg);
+
+        match *op {
+            op::Opcode::H_PREFETCH_M { rd, .. } => {
+                TimingAccess::prefetch(vec![matrix(gp(rd), matrix_prefetch)])
+            }
+            op::Opcode::H_PREFETCH_V { rd, .. } => {
+                TimingAccess::prefetch(vec![vector(gp(rd), vector_prefetch)])
+            }
+            op::Opcode::H_STORE_V { .. } => TimingAccess::Barrier,
+
+            op::Opcode::M_MM { rs1, rs2 } | op::Opcode::M_TMM { rs1, rs2 } => {
+                TimingAccess::compute(vec![
+                    matrix(gp(rs1), matrix_tile),
+                    vector(gp(rs2), matrix_vector_batch),
+                ])
+            }
+            op::Opcode::M_BMM { rs1, rs2 } | op::Opcode::M_BTMM { rs1, rs2 } => {
+                TimingAccess::compute(vec![
+                    matrix(gp(rs1), matrix_tile),
+                    vector(gp(rs2), matrix_tile),
+                ])
             }
+            op::Opcode::M_MV { rs1, rs2 } | op::Opcode::M_TMV { rs1, rs2 } => {
+                TimingAccess::compute(vec![
+                    matrix(gp(rs1), matrix_tile),
+                    vector(gp(rs2), vector_tile),
+                ])
+            }
+            op::Opcode::M_BMV { rs1, rs2, rd } | op::Opcode::M_BTMV { rs1, rs2, rd } => {
+                TimingAccess::compute(vec![
+                    matrix(gp(rs1).wrapping_add(gp(rd)), matrix_tile),
+                    vector(gp(rs2), vector_tile),
+                ])
+            }
+            op::Opcode::M_MM_WO { rd, imm, .. }
+            | op::Opcode::M_BMM_WO { rd, imm }
+            | op::Opcode::M_MV_WO { rd, imm }
+            | op::Opcode::M_BMV_WO { rd, imm } => {
+                TimingAccess::compute(vec![vector(gp(rd).wrapping_add(imm), vector_tile)])
+            }
+
+            op::Opcode::V_ADD_VV { rs1, rs2, .. }
+            | op::Opcode::V_SUB_VV { rs1, rs2, .. }
+            | op::Opcode::V_MUL_VV { rs1, rs2, .. } => TimingAccess::compute(vec![
+                vector(gp(rs1), vector_tile),
+                vector(gp(rs2), vector_tile),
+            ]),
+            op::Opcode::V_ADD_VF { rs1, .. }
+            | op::Opcode::V_SUB_VF { rs1, .. }
+            | op::Opcode::V_MUL_VF { rs1, .. }
+            | op::Opcode::V_MAX_VF { rs1, .. }
+            | op::Opcode::V_MIN_VF { rs1, .. }
+            | op::Opcode::V_EXP_V { rs1, .. }
+            | op::Opcode::V_RECI_V { rs1, .. }
+            | op::Opcode::V_RED_SUM { rs1, .. }
+            | op::Opcode::V_RED_MAX { rs1, .. }
+            | op::Opcode::V_TOPK { rs1, .. } => {
+                TimingAccess::compute(vec![vector(gp(rs1), vector_tile)])
+            }
+            op::Opcode::V_SHFT_V { rs1, .. } => {
+                TimingAccess::compute(vec![vector(gp(rs1), vector_tile)])
+            }
+
+            op::Opcode::C_BREAK => TimingAccess::Barrier,
+            _ => TimingAccess::Other,
         }
     }
 }

transactional_emulator/src/cli.rs

@@ -173,4 +173,10 @@ pub(crate) struct Opts {
     #[arg(long)]
     /// Optional JSON output path for runtime stage profile results.
     pub(crate) stage_profile_out: Option<PathBuf>,
+
+    #[arg(long, help_heading = "Experimental Timing")]
+    /// Off-by-default exploratory timing overlay that hides independent HBM
+    /// prefetch cycles behind later matrix/vector compute. This changes only
+    /// reported simulation cycles, not functional execution or HBM traffic.
+    pub(crate) experimental_overlap_prefetch_compute: bool,
 }

transactional_emulator/src/main.rs

@@ -8,9 +8,10 @@ mod op;
 mod runner;
 mod runtime_config;
 mod stage_profile;
+mod timing_overlay;
 mod vector_machine;
 
-use runtime::{Executor, Instant};
+use runtime::{Duration, Executor, Instant};
 
 #[macro_export]
 macro_rules! cycle {
@@ -30,6 +31,21 @@ async fn main() {
     let cycles = latency
         .as_picos()
         .div_ceil(runtime_config::PERIOD.as_picos().max(1));
+    if let Some(adjusted_cycles) = timing_overlay::experimental_report_cycles() {
+        let adjusted_latency =
+            Duration::from_picos(runtime_config::PERIOD.as_picos() * adjusted_cycles);
+        tracing::info!(
+            "Experimental overlap serial latency {:?} cycles {}",
+            executor.now(),
+            cycles
+        );
+        tracing::info!(
+            "Simulation completed. Latency {:?} cycles {}",
+            adjusted_latency,
+            adjusted_cycles
+        );
+        return;
+    }
     tracing::info!(
         "Simulation completed. Latency {:?} cycles {}",
         executor.now(),

transactional_emulator/src/runner.rs

@@ -16,6 +16,7 @@ use crate::runtime_config::{
     VECTOR_SRAM_SIZE, VECTOR_SRAM_TYPE, VLEN,
 };
 use crate::stage_profile::StageProfiler;
+use crate::timing_overlay::{self as timing_overlay_state, TimingOverlay};
 use crate::vector_machine::VectorMachine;
 use crate::{cli, op};
 
@@ -33,6 +34,7 @@ fn dump_to_file(path: &str, bytes: &[u8]) {
 
 pub(crate) async fn run_from_cli() {
     let opts = Opts::parse();
+    timing_overlay_state::clear_experimental_report_cycles();
 
     // If --settings is given, set PLENA_SETTINGS_TOML env var BEFORE any
     // LazyLock access (which triggers load_config()). This ensures the
@@ -215,12 +217,36 @@ pub(crate) async fn run_from_cli() {
             panic!("failed to build stage profile from ASM {:?}: {err}", path)
         })
     });
+    let mut timing_overlay = opts
+        .experimental_overlap_prefetch_compute
+        .then(TimingOverlay::default);
     accelerator
-        .do_ops(&decoded_ops, stage_profiler.as_mut())
+        .do_ops(
+            &decoded_ops,
+            stage_profiler.as_mut(),
+            timing_overlay.as_mut(),
+        )
         .await;
 
+    let serial_duration = Executor::current().now() - Instant::INIT;
+    let serial_cycles = StageProfiler::duration_to_cycles(serial_duration);
+    if let Some(overlay) = timing_overlay.as_ref() {
+        let summary = overlay.summary(serial_cycles);
+        timing_overlay_state::set_experimental_report_cycles(summary.adjusted_cycles);
+        tracing::info!(
+            serial_cycles = summary.serial_cycles,
+            adjusted_cycles = summary.adjusted_cycles,
+            hidden_prefetch_cycles = summary.hidden_prefetch_cycles,
+            pending_prefetch_cycles = summary.pending_prefetch_cycles,
+            prefetch_ops = summary.prefetch_ops,
+            compute_ops = summary.compute_ops,
+            dependent_prefetch_stalls = summary.dependent_prefetch_stalls,
+            "Experimental overlap prefetch/compute timing overlay"
+        );
+    }
+
     if let Some(profile) = stage_profiler.as_mut() {
-        profile.set_total_simulation_duration(Executor::current().now() - Instant::INIT);
+        profile.set_total_simulation_duration(serial_duration);
     }
 
     if let Some(profile) = stage_profiler.as_ref() {