use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use dashmap::DashMap;
use lru::LruCache;
use parking_lot::Mutex;
use tensor_wasm_artifacts::{ArtifactError, ArtifactStore, ContentHash, DiskArtifactStore};
use tensor_wasm_core::types::TenantId;
use zeroize::Zeroizing;
use crate::ir::TensorWasmKernelBlueprint;
use crate::ptx_emit::EmittedPtx;
#[cfg(feature = "kernel-registry")]
use crate::registry::{BlueprintResolver, KernelRegistry};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct CacheKey {
pub tenant_id: u64,
pub blueprint: u64,
pub sm_version: u32,
pub emit_config_hash: u64,
}
impl CacheKey {
pub fn for_tenant(tenant_id: TenantId, blueprint: u64, sm_version: u32) -> Self {
Self {
tenant_id: tenant_id.get(),
blueprint,
sm_version,
emit_config_hash: 0,
}
}
pub fn for_tenant_with_emit_config(
tenant_id: TenantId,
blueprint: u64,
sm_version: u32,
cfg: &crate::ptx_emit::EmitConfig,
) -> Self {
let emit_config_hash = blake3::Hasher::new()
.update(b"tensor-wasm-jit::EmitConfig::v1\0")
.update(cfg.target.as_bytes())
.update(b"\0")
.update(cfg.ptx_version.as_bytes())
.update(b"\0")
.update(&[u8::from(cfg.launch_bounds)])
.finalize();
let bytes = emit_config_hash.as_bytes();
let mut buf = [0u8; 8];
buf.copy_from_slice(&bytes[..8]);
Self {
tenant_id: tenant_id.get(),
blueprint,
sm_version,
emit_config_hash: u64::from_le_bytes(buf),
}
}
}
pub const DEFAULT_CAPACITY: usize = 256;
#[derive(Clone)]
pub struct KernelCacheConfig {
pub capacity: usize,
pub max_total_bytes: Option<u64>,
pub verify_on_get: bool,
#[cfg(feature = "kernel-registry")]
pub registry: Option<Arc<dyn KernelRegistry>>,
}
impl Default for KernelCacheConfig {
fn default() -> Self {
Self {
capacity: DEFAULT_CAPACITY,
max_total_bytes: None,
verify_on_get: true,
#[cfg(feature = "kernel-registry")]
registry: None,
}
}
}
impl KernelCacheConfig {
#[must_use]
pub fn with_capacity(mut self, capacity: usize) -> Self {
self.capacity = capacity;
self
}
#[must_use]
pub fn with_max_total_bytes(mut self, max_total_bytes: u64) -> Self {
self.max_total_bytes = Some(max_total_bytes);
self
}
#[must_use]
pub fn with_verify_on_get(mut self, on: bool) -> Self {
self.verify_on_get = on;
self
}
#[cfg(feature = "kernel-registry")]
#[must_use]
pub fn with_registry(mut self, reg: Arc<dyn KernelRegistry>) -> Self {
self.registry = Some(reg);
self
}
}
impl std::fmt::Debug for KernelCacheConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let mut d = f.debug_struct("KernelCacheConfig");
d.field("capacity", &self.capacity);
d.field("max_total_bytes", &self.max_total_bytes);
d.field("verify_on_get", &self.verify_on_get);
#[cfg(feature = "kernel-registry")]
d.field(
"registry",
&self.registry.as_ref().map(|_| "<dyn KernelRegistry>"),
);
d.finish()
}
}
#[derive(Debug, Clone)]
pub struct CachedKernel {
pub fingerprint: u64,
pub ptx: Arc<EmittedPtx>,
pub compiled: CompiledHandle,
#[doc(hidden)]
pub integrity_hash: [u8; 32],
}
impl CachedKernel {
#[must_use]
pub fn integrity_hash(&self) -> &[u8; 32] {
&self.integrity_hash
}
pub fn new(fingerprint: u64, ptx: Arc<EmittedPtx>, compiled: CompiledHandle) -> Self {
let h = blake3::hash(ptx.text.as_bytes());
Self {
fingerprint,
ptx,
compiled,
integrity_hash: *h.as_bytes(),
}
}
#[must_use]
pub fn verify_integrity(&self) -> bool {
let h = blake3::hash(self.ptx.text.as_bytes());
h.as_bytes() == &self.integrity_hash
}
}
#[derive(Debug, Clone, Default)]
pub struct CompiledHandle {
#[allow(dead_code)]
private: (),
}
#[derive(Clone)]
pub struct KernelCache {
storage: Arc<DashMap<CacheKey, Arc<CachedKernel>>>,
lru: Arc<Mutex<LruCache<CacheKey, ()>>>,
config: KernelCacheConfig,
disk: Option<Arc<DiskCache>>,
total_bytes: Arc<AtomicU64>,
verify_skipped_total: Arc<AtomicU64>,
cache_hits_total: Arc<AtomicU64>,
cache_misses_total: Arc<AtomicU64>,
integrity_reject_total: Arc<AtomicU64>,
}
impl KernelCache {
pub fn new() -> Self {
Self::with_config(KernelCacheConfig::default())
}
pub fn with_capacity(cap: usize) -> Self {
Self::with_config(KernelCacheConfig::default().with_capacity(cap))
}
pub fn with_config(mut config: KernelCacheConfig) -> Self {
config.capacity = config.capacity.max(1);
let nz = NonZeroUsize::new(config.capacity).expect(">0 (clamped above)");
Self {
storage: Arc::new(DashMap::with_capacity(config.capacity)),
lru: Arc::new(Mutex::new(LruCache::new(nz))),
config,
disk: None,
total_bytes: Arc::new(AtomicU64::new(0)),
verify_skipped_total: Arc::new(AtomicU64::new(0)),
cache_hits_total: Arc::new(AtomicU64::new(0)),
cache_misses_total: Arc::new(AtomicU64::new(0)),
integrity_reject_total: Arc::new(AtomicU64::new(0)),
}
}
#[must_use]
pub fn with_disk_persistence(mut self, cfg: DiskCacheConfig) -> Self {
self.disk = Some(Arc::new(DiskCache::new(
cfg,
Arc::clone(&self.integrity_reject_total),
)));
self
}
fn entry_bytes(kernel: &CachedKernel) -> u64 {
kernel.ptx.text.len() as u64
}
fn evict_key(&self, key: &CacheKey) -> Option<Arc<CachedKernel>> {
if let Some((_, removed)) = self.storage.remove(key) {
self.total_bytes
.fetch_sub(Self::entry_bytes(&removed), Ordering::Relaxed);
Some(removed)
} else {
None
}
}
fn enforce_byte_cap(&self) {
if let Some(cap) = self.config.max_total_bytes {
if self.total_bytes.load(Ordering::Relaxed) > cap {
let mut lru = self.lru.lock();
while self.total_bytes.load(Ordering::Relaxed) > cap && self.storage.len() > 1 {
match lru.pop_lru() {
Some((evict_key, ())) => {
self.evict_key(&evict_key);
}
None => break,
}
}
}
}
}
pub fn put(&self, key: CacheKey, kernel: CachedKernel) {
if !kernel.verify_integrity() {
tracing::error!(
target: "tensor_wasm_jit::cache",
fingerprint = kernel.fingerprint,
tenant = key.tenant_id,
"refusing to cache kernel whose integrity hash does not match \
its PTX text -- likely a struct-literal construction with a \
stale hash; use CachedKernel::new"
);
return;
}
if let Some(disk) = &self.disk {
if let Err(e) = disk.put(&key, &kernel) {
tracing::warn!(
target: "tensor_wasm_jit::cache",
fingerprint = kernel.fingerprint,
error = %e,
"disk-cache put failed; entry remains in L1 only"
);
}
}
let new_bytes = Self::entry_bytes(&kernel);
let replaced = self.storage.insert(key, Arc::new(kernel));
self.total_bytes.fetch_add(new_bytes, Ordering::Relaxed);
if let Some(old) = replaced {
self.total_bytes
.fetch_sub(Self::entry_bytes(&old), Ordering::Relaxed);
}
let evicted = self.lru.lock().push(key, ());
if let Some((evicted_key, ())) = evicted {
if evicted_key != key {
self.evict_key(&evicted_key);
}
}
if self.storage.len() > self.config.capacity {
let mut lru = self.lru.lock();
while self.storage.len() > self.config.capacity {
match lru.pop_lru() {
Some((evict_key, ())) => {
self.evict_key(&evict_key);
}
None => {
tracing::error!(
target: "tensor_wasm_jit::cache",
storage_len = self.storage.len(),
capacity = self.config.capacity,
"cache storage exceeds capacity but eviction queue is empty"
);
break;
}
}
}
}
self.enforce_byte_cap();
}
pub fn get(&self, key: &CacheKey) -> Option<Arc<CachedKernel>> {
if let Some(mut lru) = self.lru.try_lock() {
let _ = lru.get(key);
}
if let Some(entry) = self.storage.get(key) {
let kernel: Arc<CachedKernel> = Arc::clone(entry.value());
drop(entry); if self.config.verify_on_get {
if !kernel.verify_integrity() {
tracing::error!(
target: "tensor_wasm_jit::cache",
fingerprint = kernel.fingerprint,
tenant = key.tenant_id,
"L1 cache entry failed integrity verification on get; \
evicting and refusing to return it"
);
self.evict_key(key);
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
self.cache_misses_total.fetch_add(1, Ordering::Relaxed);
return None;
}
} else {
self.verify_skipped_total.fetch_add(1, Ordering::Relaxed);
if kernel.integrity_hash == [0u8; 32] {
tracing::error!(
target: "tensor_wasm_jit::cache",
fingerprint = kernel.fingerprint,
tenant = key.tenant_id,
"L1 cache entry has zero integrity_hash on verify-skip get; \
likely a struct-literal CachedKernel built without \
CachedKernel::new — evicting and refusing to return it"
);
self.evict_key(key);
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
self.cache_misses_total.fetch_add(1, Ordering::Relaxed);
return None;
}
}
self.cache_hits_total.fetch_add(1, Ordering::Relaxed);
return Some(kernel);
}
if let Some(disk) = &self.disk {
match disk.get(key) {
Ok(Some(kernel)) => {
let arc = Arc::new(kernel);
let promoted_bytes = Self::entry_bytes(&arc);
let replaced = self.storage.insert(*key, Arc::clone(&arc));
self.total_bytes
.fetch_add(promoted_bytes, Ordering::Relaxed);
if let Some(old) = replaced {
self.total_bytes
.fetch_sub(Self::entry_bytes(&old), Ordering::Relaxed);
}
if let Some((evicted_key, ())) = self.lru.lock().push(*key, ()) {
if evicted_key != *key {
self.evict_key(&evicted_key);
}
}
self.enforce_byte_cap();
self.cache_hits_total.fetch_add(1, Ordering::Relaxed);
return Some(arc);
}
Ok(None) => {}
Err(e) => {
tracing::warn!(
target: "tensor_wasm_jit::cache",
tenant = key.tenant_id,
error = %e,
"disk-cache get failed; treating as miss"
);
}
}
}
self.cache_misses_total.fetch_add(1, Ordering::Relaxed);
None
}
pub fn get_for(
&self,
tenant_id: TenantId,
blueprint: &TensorWasmKernelBlueprint,
sm_version: u32,
) -> Option<Arc<CachedKernel>> {
self.get(&CacheKey::for_tenant(
tenant_id,
blueprint.fingerprint(),
sm_version,
))
}
#[cfg(feature = "kernel-registry")]
pub fn get_with_registry_fallback(
&self,
key: &CacheKey,
resolver: &dyn BlueprintResolver,
) -> Option<Arc<CachedKernel>> {
if let Some(hit) = self.get(key) {
return Some(hit);
}
let registry = self.config.registry.as_ref()?;
let (name, version) = resolver.resolve(key.blueprint, key.sm_version)?;
let entry = registry.get(&name, &version).ok()?;
let emitted = Arc::new(crate::ptx_emit::EmittedPtx {
text: entry.1.clone(),
launch_geometry: entry.0.launch_geometry.unwrap_or((0, 0)),
});
let cached = CachedKernel::new(key.blueprint, emitted, CompiledHandle::default());
self.put(*key, cached.clone());
Some(Arc::new(cached))
}
pub fn len(&self) -> usize {
self.storage.len()
}
pub fn is_empty(&self) -> bool {
self.storage.is_empty()
}
pub fn capacity(&self) -> usize {
self.config.capacity
}
pub fn total_bytes(&self) -> u64 {
self.total_bytes.load(Ordering::Relaxed)
}
pub fn max_total_bytes(&self) -> Option<u64> {
self.config.max_total_bytes
}
pub fn active_disk_key_fingerprint(&self) -> Option<KeyFingerprint> {
self.disk.as_ref().map(|d| d.active_key_fingerprint())
}
pub fn disk_key_fingerprints(&self) -> std::io::Result<Vec<KeyFingerprint>> {
match &self.disk {
Some(d) => d.key_fingerprints_on_disk(),
None => Ok(Vec::new()),
}
}
pub fn gc_disk(&self, retain: &[KeyFingerprint]) -> std::io::Result<usize> {
match &self.disk {
Some(d) => d.gc(retain),
None => Ok(0),
}
}
pub fn config(&self) -> &KernelCacheConfig {
&self.config
}
pub fn verify_skipped_total(&self) -> u64 {
self.verify_skipped_total.load(Ordering::Relaxed)
}
pub fn cache_hits_total(&self) -> u64 {
self.cache_hits_total.load(Ordering::Relaxed)
}
pub fn cache_misses_total(&self) -> u64 {
self.cache_misses_total.load(Ordering::Relaxed)
}
pub fn integrity_reject_total(&self) -> u64 {
self.integrity_reject_total.load(Ordering::Relaxed)
}
#[doc(hidden)]
#[cfg(feature = "__unstable-test-internals")]
pub fn __test_only_insert_unchecked(&self, key: CacheKey, kernel: CachedKernel) {
let bytes = Self::entry_bytes(&kernel);
let replaced = self.storage.insert(key, Arc::new(kernel));
self.total_bytes.fetch_add(bytes, Ordering::Relaxed);
if let Some(old) = replaced {
self.total_bytes
.fetch_sub(Self::entry_bytes(&old), Ordering::Relaxed);
}
let _ = self.lru.lock().push(key, ());
}
}
#[derive(Clone)]
pub struct DiskCacheConfig {
pub dir: PathBuf,
pub hmac_key: [u8; 32],
}
impl std::fmt::Debug for DiskCacheConfig {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DiskCacheConfig")
.field("dir", &self.dir)
.field("hmac_key", &"<redacted 32 bytes>")
.finish()
}
}
impl Drop for DiskCacheConfig {
fn drop(&mut self) {
use zeroize::Zeroize;
self.hmac_key.zeroize();
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct KeyFingerprint(pub String);
impl KeyFingerprint {
#[must_use]
pub fn of_key(key: &[u8; 32]) -> Self {
let h = blake3::hash(&key[..]);
let mut buf = [0u8; 16];
hex::encode_to_slice(&h.as_bytes()[..8], &mut buf).expect("16 byte buf for 8 byte input");
Self(std::str::from_utf8(&buf).expect("hex is utf8").to_string())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl std::fmt::Display for KeyFingerprint {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.0)
}
}
struct DiskCache {
dir: PathBuf,
hmac_key: Zeroizing<[u8; 32]>,
key_prefix_hex: String,
store: Arc<DiskArtifactStore>,
integrity_reject_total: Arc<AtomicU64>,
}
const DISK_CACHE_MAGIC_V2: &[u8; 16] = b"TWJIT-KRNL-v2\0\0\0";
const DISK_CACHE_HEADER_LEN_V2: usize = 16 + 8 + 8 + 4 + 4 + 4 + 8;
const SIDECAR_MAGIC_V1: &[u8; 16] = b"TWJIT-IDX-v1\0\0\0\0";
const SIDECAR_LEN_V1: usize = 16 + 32;
impl DiskCache {
fn new(mut cfg: DiskCacheConfig, integrity_reject_total: Arc<AtomicU64>) -> Self {
let dir = std::mem::take(&mut cfg.dir);
let key_bytes = Zeroizing::new(std::mem::take(&mut cfg.hmac_key));
let store = Arc::new(DiskArtifactStore::new(dir.clone(), *key_bytes));
let key_fp = blake3::hash(&key_bytes[..]);
let mut key_prefix_buf = [0u8; 16];
hex::encode_to_slice(&key_fp.as_bytes()[..8], &mut key_prefix_buf)
.expect("16 byte buf for 8 byte input");
let key_prefix_hex = std::str::from_utf8(&key_prefix_buf)
.expect("hex is utf8")
.to_string();
Self {
dir,
hmac_key: key_bytes,
key_prefix_hex,
store,
integrity_reject_total,
}
}
fn path_for(&self, key: &CacheKey) -> PathBuf {
let mut hasher = blake3::Hasher::new();
hasher.update(b"tensor-wasm-jit::DiskCache::path::v1\0");
hasher.update(&key.tenant_id.to_le_bytes());
hasher.update(&key.blueprint.to_le_bytes());
hasher.update(&key.sm_version.to_le_bytes());
hasher.update(&key.emit_config_hash.to_le_bytes());
let h = hasher.finalize();
let digest = h.as_bytes();
let mut cache_key_hex_buf = [0u8; 32];
hex::encode_to_slice(&digest[..16], &mut cache_key_hex_buf)
.expect("32 byte buf for 16 byte input");
let cache_key_hex = std::str::from_utf8(&cache_key_hex_buf).expect("hex is utf8");
let key_prefix_hex = &self.key_prefix_hex;
self.dir
.join(format!("{key_prefix_hex}-{cache_key_hex}.ptxbin"))
}
fn encode_v2_envelope(key: &CacheKey, kernel: &CachedKernel) -> Vec<u8> {
let ptx_bytes = kernel.ptx.text.as_bytes();
let (grid_x, block_x) = kernel.ptx.launch_geometry;
let mut buf = Vec::with_capacity(DISK_CACHE_HEADER_LEN_V2 + ptx_bytes.len());
buf.extend_from_slice(DISK_CACHE_MAGIC_V2);
buf.extend_from_slice(&key.tenant_id.to_le_bytes());
buf.extend_from_slice(&key.blueprint.to_le_bytes());
buf.extend_from_slice(&key.sm_version.to_le_bytes());
buf.extend_from_slice(&grid_x.to_le_bytes());
buf.extend_from_slice(&block_x.to_le_bytes());
buf.extend_from_slice(&(ptx_bytes.len() as u64).to_le_bytes());
buf.extend_from_slice(ptx_bytes);
buf
}
fn put(&self, key: &CacheKey, kernel: &CachedKernel) -> std::io::Result<()> {
use std::io::Write;
std::fs::create_dir_all(&self.dir)?;
let envelope = Self::encode_v2_envelope(key, kernel);
let hash = self.store.put(&envelope).map_err(|e| {
std::io::Error::other(e.to_string())
})?;
let mut sidecar = Vec::with_capacity(SIDECAR_LEN_V1);
sidecar.extend_from_slice(SIDECAR_MAGIC_V1);
sidecar.extend_from_slice(hash.as_bytes());
debug_assert_eq!(sidecar.len(), SIDECAR_LEN_V1);
let sidecar_path = self.path_for(key);
let mut tmp = tempfile::NamedTempFile::new_in(&self.dir)?;
tmp.as_file_mut().write_all(&sidecar)?;
tmp.persist(&sidecar_path).map_err(std::io::Error::other)?;
Ok(())
}
fn get(&self, key: &CacheKey) -> std::io::Result<Option<CachedKernel>> {
let sidecar_path = self.path_for(key);
let sidecar = match std::fs::read(&sidecar_path) {
Ok(b) => b,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(e) => return Err(e),
};
if sidecar.len() != SIDECAR_LEN_V1 {
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
len = sidecar.len(),
"disk-cache sidecar wrong length; treating as miss"
);
return Ok(None);
}
if &sidecar[..16] != SIDECAR_MAGIC_V1 {
tracing::info!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
"cache.l2.miss.legacy_or_unknown_magic: sidecar will be rewritten on next put"
);
return Ok(None);
}
let mut hash_bytes = [0u8; 32];
hash_bytes.copy_from_slice(&sidecar[16..48]);
let content_hash = ContentHash::from_bytes(hash_bytes);
let envelope = match self.store.get(&content_hash) {
Ok(bytes) => bytes,
Err(ArtifactError::NotFound(_)) => {
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
"disk-cache sidecar references missing artifact blob; treating as miss"
);
return Ok(None);
}
Err(e) => {
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
error = %e,
"disk-cache artifact-store read failed; treating as miss"
);
return Ok(None);
}
};
if envelope.len() < DISK_CACHE_HEADER_LEN_V2 {
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
len = envelope.len(),
"disk-cache V2 envelope too short; treating as miss"
);
return Ok(None);
}
if &envelope[..16] != DISK_CACHE_MAGIC_V2 {
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
"disk-cache V2 envelope magic mismatch; treating as miss"
);
return Ok(None);
}
let mut tenant_bytes = [0u8; 8];
tenant_bytes.copy_from_slice(&envelope[16..24]);
let mut bp_bytes = [0u8; 8];
bp_bytes.copy_from_slice(&envelope[24..32]);
let mut sm_bytes = [0u8; 4];
sm_bytes.copy_from_slice(&envelope[32..36]);
let mut grid_x_bytes = [0u8; 4];
grid_x_bytes.copy_from_slice(&envelope[36..40]);
let mut block_x_bytes = [0u8; 4];
block_x_bytes.copy_from_slice(&envelope[40..44]);
let mut len_bytes = [0u8; 8];
len_bytes.copy_from_slice(&envelope[44..52]);
let tenant_on_disk = u64::from_le_bytes(tenant_bytes);
let fingerprint_on_disk = u64::from_le_bytes(bp_bytes);
let sm_version_on_disk = u32::from_le_bytes(sm_bytes);
let grid_x_on_disk = u32::from_le_bytes(grid_x_bytes);
let block_x_on_disk = u32::from_le_bytes(block_x_bytes);
let ptx_len_on_disk = u64::from_le_bytes(len_bytes) as usize;
if tenant_on_disk != key.tenant_id
|| fingerprint_on_disk != key.blueprint
|| sm_version_on_disk != key.sm_version
{
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
tenant = key.tenant_id,
tenant_on_disk,
"disk-cache V2 envelope header key mismatch (tenant/fingerprint/sm); treating as miss"
);
return Ok(None);
}
let ptx_start = DISK_CACHE_HEADER_LEN_V2;
let ptx_end = ptx_start.saturating_add(ptx_len_on_disk);
if ptx_end > envelope.len() {
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
"disk-cache declared ptx_len overruns envelope; treating as miss"
);
return Ok(None);
}
let ptx_text = match std::str::from_utf8(&envelope[ptx_start..ptx_end]) {
Ok(s) => s.to_string(),
Err(_) => {
self.integrity_reject_total.fetch_add(1, Ordering::Relaxed);
tracing::warn!(
target: "tensor_wasm_jit::cache",
file = %sidecar_path.display(),
"disk-cache PTX bytes are not valid UTF-8; treating as miss"
);
return Ok(None);
}
};
let ptx = Arc::new(EmittedPtx {
text: ptx_text,
launch_geometry: (grid_x_on_disk, block_x_on_disk),
});
Ok(Some(CachedKernel::new(
fingerprint_on_disk,
ptx,
CompiledHandle::default(),
)))
}
fn active_key_fingerprint(&self) -> KeyFingerprint {
KeyFingerprint::of_key(&self.hmac_key)
}
fn fingerprint_of_filename(name: &str) -> Option<KeyFingerprint> {
if let Some(rest) = name.strip_suffix(".ptxbin") {
let fp = rest.split('-').next()?;
if fp.is_empty() || fp.len() != 16 {
return None;
}
return Some(KeyFingerprint(fp.to_string()));
}
if let Some(rest) = name.strip_suffix(".bin") {
let fp = rest.rsplit('.').next()?;
if fp.len() != 16 {
return None;
}
return Some(KeyFingerprint(fp.to_string()));
}
None
}
fn key_fingerprints_on_disk(&self) -> std::io::Result<Vec<KeyFingerprint>> {
let mut seen = std::collections::BTreeSet::new();
let rd = match std::fs::read_dir(&self.dir) {
Ok(rd) => rd,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return Ok(Vec::new());
}
Err(e) => return Err(e),
};
for entry in rd {
let entry = entry?;
if let Some(name) = entry.file_name().to_str() {
if let Some(fp) = Self::fingerprint_of_filename(name) {
seen.insert(fp);
}
}
}
Ok(seen.into_iter().collect())
}
fn gc(&self, retain: &[KeyFingerprint]) -> std::io::Result<usize> {
let active = self.active_key_fingerprint();
let rd = match std::fs::read_dir(&self.dir) {
Ok(rd) => rd,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
return Ok(0);
}
Err(e) => return Err(e),
};
let mut removed = 0usize;
for entry in rd {
let entry = entry?;
let name_os = entry.file_name();
let name = match name_os.to_str() {
Some(n) => n,
None => continue,
};
let fp = match Self::fingerprint_of_filename(name) {
Some(fp) => fp,
None => continue, };
if fp == active || retain.contains(&fp) {
continue;
}
match std::fs::remove_file(entry.path()) {
Ok(()) => removed += 1,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
}
Err(e) => return Err(e),
}
}
Ok(removed)
}
}
impl Default for KernelCache {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::{ElemType, TensorWasmKernelBlueprint, TensorWasmOp};
use crate::ptx_emit::EmittedPtx;
use std::thread;
fn dummy_kernel(fp: u64) -> CachedKernel {
CachedKernel::new(
fp,
Arc::new(EmittedPtx {
text: String::new(),
launch_geometry: (1, 1),
}),
CompiledHandle::default(),
)
}
#[test]
fn put_then_get() {
let cache = KernelCache::new();
let key = CacheKey::for_tenant(TenantId(7), 1, 80);
cache.put(key, dummy_kernel(1));
assert_eq!(cache.get(&key).unwrap().fingerprint, 1);
}
fn sized_kernel(fp: u64, bytes: usize) -> CachedKernel {
CachedKernel::new(
fp,
Arc::new(EmittedPtx {
text: "x".repeat(bytes),
launch_geometry: (1, 1),
}),
CompiledHandle::default(),
)
}
#[test]
fn byte_cap_evicts_lru_and_bounds_total_bytes() {
let cache = KernelCache::with_config(
KernelCacheConfig::default()
.with_capacity(100)
.with_max_total_bytes(250),
);
let k1 = CacheKey::for_tenant(TenantId(0), 1, 80);
let k2 = CacheKey::for_tenant(TenantId(0), 2, 80);
let k3 = CacheKey::for_tenant(TenantId(0), 3, 80);
cache.put(k1, sized_kernel(1, 100));
cache.put(k2, sized_kernel(2, 100));
assert_eq!(cache.total_bytes(), 200, "two 100-byte entries fit");
assert_eq!(cache.len(), 2);
cache.put(k3, sized_kernel(3, 100));
assert!(
cache.total_bytes() <= 250,
"byte total must stay bounded by the cap, got {}",
cache.total_bytes()
);
assert_eq!(cache.total_bytes(), 200);
assert_eq!(cache.len(), 2);
assert!(
cache.get(&k1).is_none(),
"k1 is the LRU and must be evicted"
);
assert!(cache.get(&k2).is_some(), "k2 must survive");
assert!(cache.get(&k3).is_some(), "k3 was just inserted");
}
#[test]
fn byte_cap_admits_single_oversized_entry() {
let cache = KernelCache::with_config(
KernelCacheConfig::default()
.with_capacity(100)
.with_max_total_bytes(50),
);
let small = CacheKey::for_tenant(TenantId(0), 1, 80);
let big = CacheKey::for_tenant(TenantId(0), 2, 80);
cache.put(small, sized_kernel(1, 40));
cache.put(big, sized_kernel(2, 1000));
assert_eq!(cache.len(), 1);
assert_eq!(cache.total_bytes(), 1000);
assert!(
cache.get(&big).is_some(),
"oversized entry must still serve"
);
assert!(cache.get(&small).is_none(), "smaller LRU entry evicted");
}
#[test]
fn total_bytes_tracks_replacement() {
let cache = KernelCache::new();
let key = CacheKey::for_tenant(TenantId(0), 1, 80);
cache.put(key, sized_kernel(1, 100));
assert_eq!(cache.total_bytes(), 100);
cache.put(key, sized_kernel(1, 150));
assert_eq!(cache.len(), 1);
assert_eq!(cache.total_bytes(), 150);
}
#[test]
fn byte_cap_default_is_none() {
let cache = KernelCache::new();
assert_eq!(cache.max_total_bytes(), None);
cache.put(
CacheKey::for_tenant(TenantId(0), 1, 80),
sized_kernel(1, 512),
);
assert_eq!(cache.total_bytes(), 512);
}
#[test]
fn lru_evicts_oldest() {
let cache = KernelCache::with_capacity(2);
let k1 = CacheKey::for_tenant(TenantId(0), 1, 80);
let k2 = CacheKey::for_tenant(TenantId(0), 2, 80);
let k3 = CacheKey::for_tenant(TenantId(0), 3, 80);
cache.put(k1, dummy_kernel(1));
cache.put(k2, dummy_kernel(2));
cache.put(k3, dummy_kernel(3));
assert!(cache.get(&k1).is_none(), "k1 should have been evicted");
assert!(cache.get(&k2).is_some());
assert!(cache.get(&k3).is_some());
assert_eq!(cache.len(), 2);
}
#[test]
fn lookup_by_blueprint() {
let cache = KernelCache::new();
let bp = TensorWasmKernelBlueprint::new("k").push(TensorWasmOp::VecAdd {
elem: ElemType::F32,
lanes: 4,
});
let tenant = TenantId(11);
let key = CacheKey::for_tenant(tenant, bp.fingerprint(), 80);
cache.put(key, dummy_kernel(bp.fingerprint()));
assert!(cache.get_for(tenant, &bp, 80).is_some());
assert!(
cache.get_for(tenant, &bp, 89).is_none(),
"different sm_version is a miss"
);
assert!(
cache.get_for(TenantId(12), &bp, 80).is_none(),
"different tenant is a miss — keys are tenant-scoped"
);
}
#[test]
fn element_type_is_not_cache_aliased() {
let cache = KernelCache::new();
let tenant = TenantId(7);
let f32_bp = TensorWasmKernelBlueprint::new("k").push(TensorWasmOp::VecAdd {
elem: ElemType::F32,
lanes: 4,
});
let i32_bp = TensorWasmKernelBlueprint::new("k").push(TensorWasmOp::VecAdd {
elem: ElemType::I32,
lanes: 4,
});
assert_ne!(f32_bp.fingerprint(), i32_bp.fingerprint());
let key = CacheKey::for_tenant(tenant, f32_bp.fingerprint(), 80);
cache.put(key, dummy_kernel(f32_bp.fingerprint()));
assert!(cache.get_for(tenant, &f32_bp, 80).is_some());
assert!(
cache.get_for(tenant, &i32_bp, 80).is_none(),
"i32x4.add must not alias the f32x4.add cache slot"
);
}
#[test]
fn capacity_floor_one() {
let cache = KernelCache::with_capacity(0);
assert_eq!(cache.capacity(), 1);
}
#[test]
fn empty_when_new() {
let cache = KernelCache::new();
assert!(cache.is_empty());
assert_eq!(cache.len(), 0);
}
#[test]
fn cache_hit_returns_arc_shared_ptx() {
use std::sync::Arc;
let cache = KernelCache::new();
let bp = TensorWasmKernelBlueprint::new("matmul").push(TensorWasmOp::MatMul {
m: 16,
n: 16,
k: 16,
});
let key = CacheKey::for_tenant(TenantId(3), bp.fingerprint(), 80);
let original = Arc::new(EmittedPtx {
text: "// pre-emitted".into(),
launch_geometry: (1, 128),
});
cache.put(
key,
CachedKernel::new(
bp.fingerprint(),
original.clone(),
CompiledHandle::default(),
),
);
let hit = cache.get(&key).expect("cache hit");
assert!(Arc::ptr_eq(&hit.ptx, &original));
}
#[test]
fn concurrent_get_put_dashmap_safe() {
const N_THREADS: usize = 8;
const KEYS_PER_THREAD: u64 = 32;
let cache = KernelCache::with_capacity((N_THREADS as u64 * KEYS_PER_THREAD) as usize);
let mut handles = Vec::new();
for t in 0..N_THREADS {
let cache = cache.clone();
handles.push(thread::spawn(move || {
for i in 0..KEYS_PER_THREAD {
let key =
CacheKey::for_tenant(TenantId(0), (t as u64) * KEYS_PER_THREAD + i, 80);
cache.put(key, dummy_kernel(key.blueprint));
let _ = cache.get(&key);
}
}));
}
for h in handles {
h.join().expect("worker thread panicked");
}
for t in 0..N_THREADS {
for i in 0..KEYS_PER_THREAD {
let key = CacheKey::for_tenant(TenantId(0), (t as u64) * KEYS_PER_THREAD + i, 80);
assert!(
cache.get(&key).is_some(),
"missing key after concurrent inserts: ({t}, {i})"
);
}
}
}
#[test]
fn cache_get_returns_arc_no_clone() {
let cache = KernelCache::new();
let key = CacheKey::for_tenant(TenantId(1), 0xABCD, 80);
let original_ptx = Arc::new(EmittedPtx {
text: ".visible .entry t20_arc(){}".into(),
launch_geometry: (1, 32),
});
cache.put(
key,
CachedKernel::new(0xABCD, original_ptx.clone(), CompiledHandle::default()),
);
let first = cache.get(&key).expect("first hit");
let second = cache.get(&key).expect("second hit");
assert!(
Arc::ptr_eq(&first, &second),
"cache.get must return refcount-bumped Arc handles, not cloned \
wrappers — distinct allocations indicate the T20 perf fix \
regressed (clone-on-hit reintroduced)"
);
assert!(Arc::ptr_eq(&first.ptx, &original_ptx));
}
#[test]
fn cache_hits_misses_counter() {
let cache = KernelCache::new();
assert_eq!(cache.cache_hits_total(), 0);
assert_eq!(cache.cache_misses_total(), 0);
let key = CacheKey::for_tenant(TenantId(9), 0xC0FFEE, 80);
assert!(cache.get(&key).is_none(), "fresh cache must miss");
assert_eq!(
cache.cache_misses_total(),
1,
"first miss must bump the miss counter exactly once"
);
assert_eq!(
cache.cache_hits_total(),
0,
"a miss must not bump the hit counter"
);
cache.put(key, dummy_kernel(0xC0FFEE));
assert!(cache.get(&key).is_some(), "post-put get must hit");
assert_eq!(
cache.cache_hits_total(),
1,
"first hit must bump the hit counter exactly once"
);
assert_eq!(
cache.cache_misses_total(),
1,
"a hit must not bump the miss counter"
);
assert!(cache.get(&key).is_some());
assert_eq!(cache.cache_hits_total(), 2);
assert_eq!(cache.cache_misses_total(), 1);
}
#[test]
fn disk_envelope_binds_tenant_and_rejects_cross_tenant_blob() {
let tmp = tempfile::TempDir::new().expect("tempdir");
let dir = tmp.path().to_path_buf();
let hmac_key = [0x5Au8; 32];
let cache = KernelCache::new().with_disk_persistence(DiskCacheConfig {
dir: dir.clone(),
hmac_key,
});
let disk = cache.disk.as_ref().expect("disk configured");
let key_a = CacheKey::for_tenant(TenantId(100), 0xAA, 80);
let key_b = CacheKey::for_tenant(TenantId(200), 0xAA, 80);
let kernel = CachedKernel::new(
0xAA,
Arc::new(EmittedPtx {
text: ".visible .entry tenant_bound(){}".into(),
launch_geometry: (4, 8),
}),
CompiledHandle::default(),
);
disk.put(&key_a, &kernel).expect("put A");
let hit_a = disk.get(&key_a).expect("get A ok").expect("A present");
assert_eq!(hit_a.fingerprint, 0xAA);
assert_eq!(hit_a.ptx.launch_geometry, (4, 8));
let envelope = DiskCache::encode_v2_envelope(&key_a, &kernel);
let hash = disk.store.put(&envelope).expect("store put");
let mut sidecar = Vec::with_capacity(SIDECAR_LEN_V1);
sidecar.extend_from_slice(SIDECAR_MAGIC_V1);
sidecar.extend_from_slice(hash.as_bytes());
std::fs::write(disk.path_for(&key_b), &sidecar).expect("plant sidecar");
let before = cache.integrity_reject_total();
assert!(
disk.get(&key_b).expect("get B ok").is_none(),
"a blob bound to tenant A must not be served to tenant B"
);
assert_eq!(
cache.integrity_reject_total(),
before + 1,
"cross-tenant blob must count as an integrity rejection"
);
}
#[test]
fn disk_cache_config_debug_redacts_hmac_key() {
let cfg = DiskCacheConfig {
dir: PathBuf::from("/tmp/tensor-wasm-jit-debug-test"),
hmac_key: [0xDEu8; 32],
};
let dbg = format!("{cfg:?}");
let lowered = dbg.to_ascii_lowercase();
let decimal_hits = lowered.matches("222").count();
let hex_hits = lowered.matches("de").count();
assert!(
decimal_hits < 32,
"DiskCacheConfig Debug output appears to contain the raw key in \
decimal form ({decimal_hits} occurrences of \"222\"): {dbg}"
);
assert!(
hex_hits < 32,
"DiskCacheConfig Debug output appears to contain the raw key in \
hex form ({hex_hits} occurrences of \"de\"): {dbg}"
);
assert!(
lowered.contains("redacted"),
"DiskCacheConfig Debug output should mark hmac_key as redacted: {dbg}"
);
}
}