sim_lib_compute_wgpu/
pipeline.rs

1//! Bounded validated pipeline cache for portable wgpu tensor kernels.
2
3use std::collections::VecDeque;
4use std::sync::Arc;
5
6use sim_kernel::Symbol;
7
8use crate::{WgpuAdapterProbe, kernels::kernel_wgsl_for_op};
9
10/// Portable operation implemented by a wgpu tensor kernel.
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub enum WgpuKernelOp {
13    /// Element-wise addition.
14    Add,
15    /// Element-wise subtraction.
16    Sub,
17    /// Element-wise multiplication.
18    Mul,
19    /// Element-wise division.
20    Div,
21    /// Element-wise negation.
22    Neg,
23    /// Element-wise square root.
24    Sqrt,
25    /// Element-wise exponential.
26    Exp,
27    /// Element-wise natural logarithm.
28    Log,
29    /// Element-wise sine.
30    Sin,
31    /// Element-wise cosine.
32    Cos,
33    /// Whole-tensor sum reduction.
34    Sum,
35    /// Whole-tensor minimum reduction.
36    Min,
37    /// Whole-tensor maximum reduction.
38    Max,
39    /// Whole-tensor Euclidean norm.
40    Norm,
41    /// Matrix transpose.
42    Transpose,
43    /// Vector dot product.
44    Dot,
45    /// Vector or matrix multiplication.
46    Matmul,
47}
48
49impl WgpuKernelOp {
50    /// Returns true when the kernel consumes two input tensors.
51    pub fn is_binary(self) -> bool {
52        matches!(self, Self::Add | Self::Sub | Self::Mul | Self::Div)
53    }
54
55    /// Returns true when the kernel consumes exactly one input tensor.
56    pub fn is_unary(self) -> bool {
57        matches!(
58            self,
59            Self::Sqrt
60                | Self::Neg
61                | Self::Exp
62                | Self::Log
63                | Self::Sin
64                | Self::Cos
65                | Self::Sum
66                | Self::Min
67                | Self::Max
68                | Self::Norm
69                | Self::Transpose
70        )
71    }
72
73    /// Returns true when the kernel performs fixed-tree accumulation.
74    pub fn is_reduction(self) -> bool {
75        matches!(self, Self::Sum | Self::Min | Self::Max | Self::Norm)
76    }
77
78    /// Returns true when the kernel performs a linalg memory/product primitive.
79    pub fn is_linalg(self) -> bool {
80        matches!(self, Self::Transpose | Self::Dot | Self::Matmul)
81    }
82
83    fn symbol_name(self) -> &'static str {
84        match self {
85            Self::Add => "add",
86            Self::Sub => "sub",
87            Self::Mul => "mul",
88            Self::Div => "div",
89            Self::Neg => "neg",
90            Self::Sqrt => "sqrt",
91            Self::Exp => "exp",
92            Self::Log => "log",
93            Self::Sin => "sin",
94            Self::Cos => "cos",
95            Self::Sum => "sum",
96            Self::Min => "min",
97            Self::Max => "max",
98            Self::Norm => "norm",
99            Self::Transpose => "transpose",
100            Self::Dot => "dot",
101            Self::Matmul => "matmul",
102        }
103    }
104
105    fn wgsl_bytes(self) -> usize {
106        kernel_wgsl_for_op(self).len()
107    }
108}
109
110/// A compiled pipeline paired with its stable evidence record.
111#[derive(Clone, Debug)]
112pub struct WgpuCompiledPipeline {
113    /// Public evidence for this validated pipeline.
114    pub record: WgpuPipelineRecord,
115    /// Retained native compute pipeline.
116    pub pipeline: Arc<wgpu::ComputePipeline>,
117}
118
119/// Dtype strategy selected for a portable kernel.
120#[derive(Clone, Copy, Debug, PartialEq, Eq)]
121pub enum WgpuKernelDType {
122    /// Native f32 arithmetic.
123    F32,
124    /// Native f16 arithmetic, only when the adapter granted shader f16.
125    F16Native,
126    /// bf16/unsupported half inputs widened to f32 arithmetic.
127    Bf16WidenedToF32,
128}
129
130impl WgpuKernelDType {
131    fn symbol_name(self) -> &'static str {
132        match self {
133            Self::F32 => "f32",
134            Self::F16Native => "f16",
135            Self::Bf16WidenedToF32 => "f32-widened-half",
136        }
137    }
138}
139
140/// Cache key for a validated portable pipeline.
141#[derive(Clone, Debug, PartialEq, Eq)]
142pub struct WgpuPipelineKey {
143    /// Adapter ordinal from discovery evidence.
144    pub adapter_ordinal: usize,
145    /// Portable kernel operation.
146    pub op: WgpuKernelOp,
147    /// Kernel dtype strategy.
148    pub dtype: WgpuKernelDType,
149    /// Output rank.
150    pub rank: usize,
151    /// Bounded planner profile selected from granted limits.
152    pub tile: WgpuTileProfile,
153}
154
155/// Portable tile profile selected from granted adapter limits.
156#[derive(Clone, Copy, Debug, PartialEq, Eq)]
157pub struct WgpuTileProfile {
158    /// Workgroup width for one-dimensional kernels.
159    pub workgroup_width: u32,
160    /// Square matrix tile edge used by transpose and matmul.
161    pub matrix_tile: u32,
162    /// Number of workgroup-local reduction lanes.
163    pub reduction_lanes: u32,
164    /// Maximum resident bytes accepted by one planned dispatch.
165    pub max_dispatch_bytes: u64,
166}
167
168impl WgpuTileProfile {
169    /// Selects conservative portable tiles from granted adapter limits.
170    pub fn from_probe(probe: &WgpuAdapterProbe) -> Self {
171        let limits = &probe.adapter.granted_limits;
172        let max_x = limits.max_compute_workgroup_size_x.max(1);
173        let max_invocations = limits.max_compute_invocations_per_workgroup.max(1);
174        let workgroup_width = max_x.min(max_invocations).clamp(1, 256);
175        let matrix_tile = 16_u32.min(workgroup_width).min(max_invocations).max(1);
176        let reduction_lanes = workgroup_width.clamp(1, 256);
177        Self {
178            workgroup_width,
179            matrix_tile,
180            reduction_lanes,
181            max_dispatch_bytes: limits.max_buffer_size.max(4),
182        }
183    }
184}
185
186/// Validated pipeline evidence retained in the cache.
187#[derive(Clone, Debug, PartialEq, Eq)]
188pub struct WgpuPipelineRecord {
189    /// Cache key.
190    pub key: WgpuPipelineKey,
191    /// Stable symbol for this validated pipeline.
192    pub symbol: Symbol,
193    /// WGSL source byte length.
194    pub wgsl_bytes: usize,
195}
196
197/// Snapshot of cache pressure and reuse.
198#[derive(Clone, Debug, Default, PartialEq, Eq)]
199pub struct WgpuPipelineCacheSnapshot {
200    /// Cached pipeline count.
201    pub entries: usize,
202    /// Number of cache hits.
203    pub hits: usize,
204    /// Number of cache misses.
205    pub misses: usize,
206    /// Number of oldest-entry evictions.
207    pub evictions: usize,
208}
209
210/// Small FIFO cache for validated pipelines.
211#[derive(Clone, Debug)]
212pub struct WgpuPipelineCache {
213    capacity: usize,
214    entries: VecDeque<WgpuPipelineRecord>,
215    compiled: VecDeque<(WgpuPipelineKey, Arc<wgpu::ComputePipeline>)>,
216    hits: usize,
217    misses: usize,
218    evictions: usize,
219}
220
221impl WgpuPipelineCache {
222    /// Builds a cache with a bounded entry count.
223    pub fn new(capacity: usize) -> Self {
224        Self {
225            capacity: capacity.max(1),
226            entries: VecDeque::new(),
227            compiled: VecDeque::new(),
228            hits: 0,
229            misses: 0,
230            evictions: 0,
231        }
232    }
233
234    /// Returns an existing pipeline or validates and inserts one.
235    pub fn get_or_insert(
236        &mut self,
237        probe: &WgpuAdapterProbe,
238        op: WgpuKernelOp,
239        dtype: WgpuKernelDType,
240        rank: usize,
241    ) -> WgpuPipelineRecord {
242        let tile = WgpuTileProfile::from_probe(probe);
243        let key = WgpuPipelineKey {
244            adapter_ordinal: probe.adapter.ordinal,
245            op,
246            dtype,
247            rank,
248            tile,
249        };
250        if let Some(record) = self.entries.iter().find(|record| record.key == key) {
251            self.hits += 1;
252            return record.clone();
253        }
254        self.misses += 1;
255        if self.entries.len() == self.capacity {
256            if let Some(evicted) = self.entries.pop_front() {
257                self.compiled
258                    .retain(|(compiled_key, _)| *compiled_key != evicted.key);
259            }
260            self.evictions += 1;
261        }
262        let record = WgpuPipelineRecord {
263            symbol: Symbol::qualified(
264                "compute.pipeline.wgpu",
265                format!(
266                    "{}/{}/{}/rank-{}",
267                    probe.adapter.ordinal,
268                    op.symbol_name(),
269                    dtype.symbol_name(),
270                    rank
271                ),
272            ),
273            key,
274            wgsl_bytes: op.wgsl_bytes(),
275        };
276        self.entries.push_back(record.clone());
277        record
278    }
279
280    /// Returns an existing compiled pipeline or validates, compiles, and inserts one.
281    pub fn get_or_insert_compiled(
282        &mut self,
283        device: &wgpu::Device,
284        probe: &WgpuAdapterProbe,
285        op: WgpuKernelOp,
286        dtype: WgpuKernelDType,
287        rank: usize,
288    ) -> WgpuCompiledPipeline {
289        let record = self.get_or_insert(probe, op, dtype, rank);
290        if let Some((_, pipeline)) = self
291            .compiled
292            .iter()
293            .find(|(compiled_key, _)| *compiled_key == record.key)
294        {
295            return WgpuCompiledPipeline {
296                record,
297                pipeline: pipeline.clone(),
298            };
299        }
300        let label = if op.is_reduction() {
301            "sim-compute-wgpu-reduction"
302        } else if op.is_linalg() {
303            "sim-compute-wgpu-linalg"
304        } else {
305            "sim-compute-wgpu-pointwise"
306        };
307        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
308            label: Some(label),
309            source: wgpu::ShaderSource::Wgsl(kernel_wgsl_for_op(op).into()),
310        });
311        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
312            label: Some(label),
313            layout: None,
314            module: &shader,
315            entry_point: Some("main"),
316            compilation_options: Default::default(),
317            cache: None,
318        });
319        let pipeline = Arc::new(pipeline);
320        self.compiled
321            .push_back((record.key.clone(), pipeline.clone()));
322        WgpuCompiledPipeline { record, pipeline }
323    }
324
325    /// Returns cache pressure and reuse counters.
326    pub fn snapshot(&self) -> WgpuPipelineCacheSnapshot {
327        WgpuPipelineCacheSnapshot {
328            entries: self.entries.len(),
329            hits: self.hits,
330            misses: self.misses,
331            evictions: self.evictions,
332        }
333    }
334}
335
336impl Default for WgpuPipelineCache {
337    fn default() -> Self {
338        Self::new(16)
339    }
340}
sim_lib_compute_wgpu/pipeline.rs

sim_lib_compute_wgpu/
pipeline.rs
