prefetch data into cpu cache aot

Author: Human9000-bit

crates/yscv-kernels/README.md

@@ -61,7 +61,9 @@ packing loop.
 Cascade dispatch: 48→32→16→8→4→scalar columns. k-loop unrolled by 2
 with doubled accumulator sets to break FMA latency dependency chains
 (FMA latency = 4 cycles, 2 FMA ports on Zen 4 — spacing accumulator
-reuse to 4+ cycles eliminates pipeline stalls).
+reuse to 4+ cycles eliminates pipeline stalls). Four-row transposed-A tiles
+prefetch strided B panels four K iterations ahead for the measured
+`M >= 128, K <= 32` AVX/FMA range and in-order Cortex-A53/A55 kernels.
 
 ### Depthwise Conv SIMD
 

crates/yscv-kernels/benches/kernels_cpu_ops.rs

@@ -2,12 +2,12 @@ use std::num::NonZeroUsize;
 
 use criterion::{Criterion, black_box, criterion_group, criterion_main};
 use yscv_kernels::{
-    Backend, BatchNorm2dParams, LayerNormLastDimParams, ParallelElementwiseConfig,
+    Backend, BatchNorm2dParams, BinaryKind, LayerNormLastDimParams, ParallelElementwiseConfig,
     ParallelMatmulConfig, SeparableConv2dParams, ThreadedCpuBackend, ThreadedCpuBackendConfig, add,
-    avg_pool2d_nhwc, batch_norm2d_nhwc, conv2d_nhwc, conv2d_nhwc_indirect_padded,
-    conv2d_nhwc_padded, depthwise_conv2d_nhwc, layer_norm_last_dim, log_softmax_last_dim,
-    logsumexp_last_dim, matmul_2d, matmul_2d_sequential, max_pool2d_nhwc, relu,
-    separable_conv2d_nhwc, sigmoid, softmax_last_dim,
+    avg_pool2d_nhwc, batch_norm2d_nhwc, binary_same_shape_dispatch, conv2d_nhwc,
+    conv2d_nhwc_indirect_padded, conv2d_nhwc_padded, depthwise_conv2d_nhwc, layer_norm_last_dim,
+    log_softmax_last_dim, logsumexp_last_dim, matmul_2d, matmul_2d_sequential, max_pool2d_nhwc,
+    relu, separable_conv2d_nhwc, sigmoid, softmax_last_dim,
 };
 use yscv_tensor::Tensor;
 
@@ -170,6 +170,27 @@ fn bench_elementwise_modes(c: &mut Criterion) {
             black_box(out);
         });
     });
+    let mut raw_out = vec![0.0; lhs.data().len()];
+    group.bench_function("add_same_shape_raw_slice", |b| {
+        b.iter(|| {
+            binary_same_shape_dispatch(
+                black_box(lhs.data()),
+                black_box(rhs.data()),
+                black_box(&mut raw_out),
+                BinaryKind::Add,
+            );
+        });
+    });
+    group.bench_function("mul_same_shape_raw_slice", |b| {
+        b.iter(|| {
+            binary_same_shape_dispatch(
+                black_box(lhs.data()),
+                black_box(rhs.data()),
+                black_box(&mut raw_out),
+                BinaryKind::Mul,
+            );
+        });
+    });
     group.bench_function("add_same_shape_threaded_2", |b| {
         b.iter(|| {
             let out = threaded_backend

crates/yscv-kernels/src/ops/matmul/trans_a.rs

@@ -4,6 +4,25 @@
 
 use super::*;
 
+#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
+const MATMUL_PREFETCH_AHEAD: usize = 4;
+
+#[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))]
+#[inline(always)]
+#[allow(unsafe_code, unsafe_op_in_unsafe_fn)]
+unsafe fn prefetch_l1_keep(p: *const f32) {
+    #[cfg(target_arch = "x86")]
+    std::arch::x86::_mm_prefetch::<{ std::arch::x86::_MM_HINT_T0 }>(p as *const i8);
+    #[cfg(target_arch = "x86_64")]
+    std::arch::x86_64::_mm_prefetch::<{ std::arch::x86_64::_MM_HINT_T0 }>(p as *const i8);
+    #[cfg(target_arch = "aarch64")]
+    core::arch::asm!(
+        "prfm pldl1keep, [{p}]",
+        p = in(reg) p,
+        options(nostack, preserves_flags, readonly),
+    );
+}
+
 pub(super) fn non_trans_4row_disabled() -> bool {
     static CACHED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
     *CACHED.get_or_init(|| std::env::var_os("YSCV_NON_TRANS_4ROW_OFF").is_some())
@@ -929,6 +948,8 @@ unsafe fn trans_a_4row_avx2(
     use std::arch::x86::*;
     #[cfg(target_arch = "x86_64")]
     use std::arch::x86_64::*;
+    let prefetch = m >= 128 && n > 16 && k <= 32;
+    let prefetch_end = k.saturating_sub(MATMUL_PREFETCH_AHEAD);
     unsafe {
         let r0 = out_4rows.as_mut_ptr();
         let r1 = r0.add(n);
@@ -955,6 +976,9 @@ unsafe fn trans_a_4row_avx2(
                 let a1 = _mm256_set1_ps(*a_row.add(1));
                 let a2 = _mm256_set1_ps(*a_row.add(2));
                 let a3 = _mm256_set1_ps(*a_row.add(3));
+                if prefetch && ki < prefetch_end {
+                    prefetch_l1_keep(b.as_ptr().add((ki + MATMUL_PREFETCH_AHEAD) * n + col));
+                }
                 let bptr = b.as_ptr().add(ki * n + col);
                 let b0 = _mm256_loadu_ps(bptr);
                 let b1 = _mm256_loadu_ps(bptr.add(8));
@@ -1015,6 +1039,8 @@ unsafe fn trans_a_4row_neon(
     n: usize,
     out_4rows: &mut [f32],
 ) {
+    let prefetch = n > 16 && crate::host_cpu().uarch.is_in_order();
+    let prefetch_end = k.saturating_sub(MATMUL_PREFETCH_AHEAD);
     unsafe {
         let r0 = out_4rows.as_mut_ptr();
         let r1 = r0.add(n);
@@ -1048,6 +1074,9 @@ unsafe fn trans_a_4row_neon(
                 let a1 = vdupq_n_f32(*a_row.add(1));
                 let a2 = vdupq_n_f32(*a_row.add(2));
                 let a3 = vdupq_n_f32(*a_row.add(3));
+                if prefetch && ki < prefetch_end {
+                    prefetch_l1_keep(b.as_ptr().add((ki + MATMUL_PREFETCH_AHEAD) * n + col));
+                }
                 let bptr = b.as_ptr().add(ki * n + col);
                 let b0 = vld1q_f32(bptr);
                 let b1 = vld1q_f32(bptr.add(4));

crates/yscv-kernels/src/ops/simd/binary.rs

@@ -507,23 +507,10 @@ unsafe fn binary_same_shape_avx(lhs: &[f32], rhs: &[f32], out: &mut [f32], kind:
     let out_ptr = out.as_mut_ptr();
     let mut index = 0usize;
 
-    // 4x unrolled: process 32 floats per iteration with software prefetch.
-    // Matches vDSP throughput by keeping the OoO pipeline fully saturated.
+    // 4x unrolled: process 32 floats per iteration.
     match kind {
         BinaryKind::Add => {
             while index + 32 <= len {
-                #[cfg(target_arch = "x86")]
-                {
-                    use std::arch::x86::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
-                #[cfg(target_arch = "x86_64")]
-                {
-                    use std::arch::x86_64::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
                 let a0 = _mm256_loadu_ps(left_ptr.add(index));
                 let b0 = _mm256_loadu_ps(right_ptr.add(index));
                 let a1 = _mm256_loadu_ps(left_ptr.add(index + 8));
@@ -541,18 +528,6 @@ unsafe fn binary_same_shape_avx(lhs: &[f32], rhs: &[f32], out: &mut [f32], kind:
         }
         BinaryKind::Sub => {
             while index + 32 <= len {
-                #[cfg(target_arch = "x86")]
-                {
-                    use std::arch::x86::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
-                #[cfg(target_arch = "x86_64")]
-                {
-                    use std::arch::x86_64::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
                 let a0 = _mm256_loadu_ps(left_ptr.add(index));
                 let b0 = _mm256_loadu_ps(right_ptr.add(index));
                 let a1 = _mm256_loadu_ps(left_ptr.add(index + 8));
@@ -570,18 +545,6 @@ unsafe fn binary_same_shape_avx(lhs: &[f32], rhs: &[f32], out: &mut [f32], kind:
         }
         BinaryKind::Mul => {
             while index + 32 <= len {
-                #[cfg(target_arch = "x86")]
-                {
-                    use std::arch::x86::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
-                #[cfg(target_arch = "x86_64")]
-                {
-                    use std::arch::x86_64::_mm_prefetch;
-                    _mm_prefetch::<3>(left_ptr.add(index + 32) as *const i8);
-                    _mm_prefetch::<3>(right_ptr.add(index + 32) as *const i8);
-                }
                 let a0 = _mm256_loadu_ps(left_ptr.add(index));
                 let b0 = _mm256_loadu_ps(right_ptr.add(index));
                 let a1 = _mm256_loadu_ps(left_ptr.add(index + 8));