Skip to content

Various kernel opts #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Human9000-bit wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Various kernel opts #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion crates/yscv-kernels/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
31 changes: 26 additions & 5 deletions crates/yscv-kernels/benches/kernels_cpu_ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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
Expand Down
29 changes: 29 additions & 0 deletions crates/yscv-kernels/src/ops/matmul/trans_a.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down Expand Up @@ -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);
Expand All @@ -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));
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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));
Expand Down
39 changes: 1 addition & 38 deletions crates/yscv-kernels/src/ops/simd/binary.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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));
Expand All @@ -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));
Expand Down
Loading
Loading