pictor_kernels/
lib.rs

1#![cfg_attr(
2    all(target_arch = "aarch64", nightly_aarch64_prefetch),
3    feature(stdarch_aarch64_prefetch)
4)]
5
6//! # pictor-kernels
7//!
8//! 1-bit Q1\_0\_g128 compute kernels for Pictor.
9//!
10//! Provides dequantization and fused matrix-multiply operations optimized
11//! for the PrismML 1-bit weight format. The kernels are organised in a
12//! tiered dispatch architecture that auto-selects the fastest implementation
13//! available on the current CPU:
14//!
15//! | Tier | Feature gate | Instruction set |
16//! |------|-------------|-----------------|
17//! | **Reference** | always | Pure scalar Rust (correctness baseline) |
18//! | **AVX2+FMA** | `simd-avx2` | 256-bit SIMD (x86-64) |
19//! | **AVX-512** | `simd-avx512` | 512-bit SIMD (x86-64) |
20//! | **NEON** | `simd-neon` | 128-bit SIMD (AArch64) |
21//!
22//! Runtime dispatch is handled by [`KernelDispatcher`] which queries
23//! SciRS2-Core's SIMD capability cache on construction.
24//!
25//! ## Key Kernels
26//!
27//! | Kernel | Description |
28//! |--------|-------------|
29//! | [`dequant::dequant_1bit_g128`] | Unpack 128 sign bits + FP16 scale → FP32 |
30//! | [`gemv::gemv_1bit_g128`] | 1-bit weight matrix × FP32 vector (single-token decode) |
31//! | [`gemm::gemm_1bit_g128`] | 1-bit weight matrix × FP32 matrix (multi-token prefill) |
32//!
33//! ## Trait
34//!
35//! All tiers implement [`OneBitKernel`] so callers are agnostic to the
36//! underlying SIMD level.
37
38/// Emit an AArch64 software-prefetch hint, degrading to a no-op off-nightly.
39///
40/// `core::arch::aarch64::_prefetch` is gated behind the `stdarch_aarch64_prefetch`
41/// nightly feature. On nightly AArch64 this expands to the real intrinsic
42/// (identical codegen); on stable — or any non-AArch64 target — it expands to a
43/// no-op that merely consumes `$ptr`. Prefetch is a pure performance hint, so
44/// dropping it never changes any computed result (correctness/parity unaffected).
45///
46/// `$ptr` must be `*const i8`; `$rw` (0 = read, 1 = write) and `$loc`
47/// (0..=3 cache locality) must be const expressions, matching the intrinsic ABI.
48///
49/// `allow(unused_macros)`: every invocation lives in `#[cfg(target_arch =
50/// "aarch64")]` code (`prefetch.rs`, `simd_neon.rs`), so on x86_64 / other
51/// targets the macro is defined-but-unused. It is kept defined on all targets
52/// (rather than cfg-gated away) to preserve its cross-platform no-op contract.
53#[allow(unused_macros)]
54macro_rules! aarch64_prefetch {
55    ($ptr:expr, $rw:expr, $loc:expr) => {{
56        // SAFETY: prefetch is always safe — invalid addresses are silently
57        // ignored on ARM. `$rw`/`$loc` are const as required by the intrinsic.
58        // The macro is invoked from both safe fns (where the `unsafe` block is
59        // required) and `unsafe fn` bodies (where it is redundant), so
60        // `unused_unsafe` is allowed to keep both call sites warning-free.
61        #[cfg(all(target_arch = "aarch64", nightly_aarch64_prefetch))]
62        #[allow(unused_unsafe)]
63        unsafe {
64            core::arch::aarch64::_prefetch($ptr, $rw, $loc);
65        }
66        #[cfg(not(all(target_arch = "aarch64", nightly_aarch64_prefetch)))]
67        {
68            let _ = $ptr;
69        }
70    }};
71}
72#[allow(unused_imports)] // re-export unused on non-aarch64 targets (see macro doc above)
73pub(crate) use aarch64_prefetch;
74
75#[cfg(all(feature = "metal", target_os = "macos"))]
76#[macro_use]
77extern crate objc;
78
79pub mod gpu_backend;
80#[cfg(feature = "gpu")]
81pub use gpu_backend::Scirs2Backend;
82pub use gpu_backend::{
83    gpu_gemv_1bit, gpu_matmul, select_backend, CpuBackend, DeviceBuffer, GpuBackend,
84    GpuBackendTrait, GpuError, LaunchConfig,
85};
86
87#[cfg(all(feature = "metal", target_os = "macos"))]
88pub use gpu_backend::{
89    build_cached_weights, build_cached_weights_ternary_only, metal_fused_gate_up_swiglu_fp8_e4m3,
90    metal_fused_gate_up_swiglu_fp8_e5m2, metal_gemm_fp8_e4m3, metal_gemm_fp8_e4m3_residual,
91    metal_gemm_fp8_e5m2, metal_gemm_fp8_e5m2_residual, metal_gemv_fp8_e4m3, metal_gemv_fp8_e5m2,
92    print_gpu_profile_summary, try_metal_ffn, try_metal_forward_greedy_ternary,
93    try_metal_full_forward, try_metal_full_forward_cached, try_metal_full_forward_prefill,
94    try_metal_full_forward_prefill_ternary, try_metal_full_forward_prefill_verify,
95    try_metal_full_forward_prefill_verify_ternary, try_metal_full_forward_ternary,
96    try_metal_full_layer, try_metal_prefill_ternary, try_metal_prefill_verify_ternary,
97    try_metal_qkv, CachedLayerWeights, CachedModelWeights, FullForwardLayerParams,
98    FullForwardLayerParamsTernary, MetalGraph, MetalGraphError, MetalWeightHandle,
99};
100
101#[cfg(all(
102    feature = "native-cuda",
103    any(target_os = "linux", target_os = "windows")
104))]
105pub use gpu_backend::{
106    cuda_gemv_q2k, cuda_gemv_q3k, cuda_gemv_q4k, cuda_gemv_q5k, cuda_gemv_q6k, cuda_gemv_q8k,
107};
108
109#[cfg(all(
110    feature = "native-cuda",
111    any(target_os = "linux", target_os = "windows")
112))]
113pub use gpu_backend::{
114    cuda_gemv_fp8_e4m3, cuda_gemv_fp8_e5m2, cuda_gemv_q4_0, cuda_gemv_q8_0, try_cuda_ffn,
115    try_cuda_full_forward, try_cuda_full_forward_ternary,
116    try_cuda_full_forward_ternary_with_gpu_lm_head, try_cuda_full_forward_with_gpu_lm_head,
117    try_cuda_full_layer, try_cuda_prefill, try_cuda_prefill_q_std, try_cuda_prefill_ternary,
118    try_cuda_qkv, CudaCachedLayerWeights, CudaFullForwardLayerParams,
119    CudaFullForwardLayerParamsTernary, CudaGraph, CudaGraphError, CudaQStdPrefillLayerParams,
120    DitSingleBlockWeights, NativeCudaBackend,
121};
122
123#[cfg(all(
124    feature = "native-cuda",
125    any(target_os = "linux", target_os = "windows")
126))]
127pub use gpu_backend::{try_cuda_prefill_k_quant, CudaKQuantPrefillLayerParams, KQuantFormat};
128
129#[cfg(all(
130    feature = "native-cuda",
131    any(target_os = "linux", target_os = "windows")
132))]
133pub use gpu_backend::{try_cuda_prefill_fp8, CudaFP8PrefillLayerParams};
134
135pub mod dequant;
136pub mod dequant_fp8;
137pub mod dequant_ternary;
138pub mod dispatch;
139pub mod error;
140pub mod fp8_lut;
141pub mod gemm;
142pub mod gemm_fp8;
143pub mod gemm_ternary;
144pub mod gemv;
145pub mod gemv_fp8;
146pub mod gemv_q2k;
147pub mod gemv_q3k;
148pub mod gemv_q4_0;
149pub mod gemv_q4k;
150pub mod gemv_q5k;
151pub mod gemv_q6k;
152pub mod gemv_q8_0;
153pub mod gemv_q8k;
154pub mod gemv_ternary;
155pub mod packing;
156pub mod parallel;
157pub mod parallel_tiled;
158#[cfg(target_arch = "x86_64")]
159pub mod simd_avx2;
160#[cfg(target_arch = "x86_64")]
161pub mod simd_avx512;
162#[cfg(target_arch = "x86_64")]
163pub mod simd_fp8_avx2;
164#[cfg(target_arch = "x86_64")]
165pub mod simd_fp8_avx512;
166#[cfg(target_arch = "aarch64")]
167pub mod simd_fp8_neon;
168#[cfg(target_arch = "aarch64")]
169pub mod simd_neon;
170pub mod tiled;
171pub mod traits;
172pub mod weight_cache;
173
174pub mod aligned;
175pub mod prefetch;
176pub mod simd_float_ops;
177pub mod tuning;
178
179pub use aligned::{AlignedBlocks, AlignedBuffer};
180pub use dispatch::{KernelDispatcher, KernelTier};
181pub use error::{KernelError, KernelResult};
182pub use gemv_q2k::gemv_q2k;
183pub use gemv_q3k::gemv_q3k;
184pub use gemv_q4_0::gemv_q4_0;
185pub use gemv_q4k::gemv_q4k;
186pub use gemv_q5k::gemv_q5k;
187pub use gemv_q6k::gemv_q6k;
188pub use gemv_q8_0::gemv_q8_0;
189pub use gemv_q8k::gemv_q8k;
190pub use parallel::{
191    gemm_fp8_e4m3_par, gemm_fp8_e5m2_par, gemm_ternary_g128_par, gemv_fp8_e4m3_par,
192    gemv_fp8_e5m2_par, gemv_ternary_g128_par,
193};
194pub use parallel_tiled::{gemm_adaptive_ternary, gemv_adaptive, gemv_adaptive_ternary};
195pub use prefetch::{PrefetchConfig, PrefetchLocality, PrefetchStrategy};
196pub use simd_float_ops::{rms_norm_simd, rope_apply_simd, silu_simd, softmax_simd, swiglu_simd};
197pub use traits::{Fp8Kernel, OneBitKernel, TernaryKernel};
198pub use tuning::{PlatformProfile, TunedThresholds, TuningSummary};
199pub use weight_cache::GpuWeightHandle;
pictor_kernels/lib.rs

pictor_kernels/
lib.rs
