rusolver/
kernel_plan.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Dependency-free launch geometry for the opt-in warp-shared direct solvers.
3//! This describes source-level storage/work, NOT measured bandwidth or latency.
4use core::fmt;
5
6/// Which scratch layout is needed. Both algorithms preserve the serial APIs.
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub enum WarpDirectKind { Cholesky, Lu }
9
10/// Immutable validated plan; one 32-thread block owns one system.
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
12pub struct WarpDirectPlan {
13    batch: usize,
14    order: usize,
15    rhs: usize,
16    pitch: usize,
17    shared_bytes: usize,
18    kind: WarpDirectKind,
19}
20
21/// Capability snapshot supplied by a real runtime, not inferred from its name.
22#[derive(Clone, Copy, Debug)]
23pub struct WarpDirectLimits {
24    pub plane_min: u32,
25    pub plane_max: u32,
26    pub plane_ops: bool,
27    pub max_threads: u32,
28    pub max_block_x: u32,
29    pub max_grid_x: u32,
30    pub shared_bytes: usize,
31}
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq)]
34pub enum WarpPlanError { InvalidShape, SizeOverflow, Unsupported(&'static str) }
35impl fmt::Display for WarpPlanError {
36    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
37        match self {
38            Self::InvalidShape => f.write_str("require n=1..32 and nrhs=1..8"),
39            Self::SizeOverflow => f.write_str("32-bit device indexing/byte count overflow"),
40            Self::Unsupported(s) => write!(f, "warp solver requires {s}"),
41        }
42    }
43}
44impl std::error::Error for WarpPlanError {}
45
46impl WarpDirectPlan {
47    pub const THREADS: u32 = 32;
48    pub fn new(batch: usize, order: usize, rhs: usize, kind: WarpDirectKind) -> Result<Self, WarpPlanError> {
49        if !(1..=32).contains(&order) || !(1..=8).contains(&rhs) {
50            return Err(WarpPlanError::InvalidShape);
51        }
52        for per_system in [order * order, order * rhs, order, 1] {
53            let elements = batch.checked_mul(per_system).ok_or(WarpPlanError::SizeOverflow)?;
54            if elements > u32::MAX as usize { return Err(WarpPlanError::SizeOverflow); }
55            elements.checked_mul(4).ok_or(WarpPlanError::SizeOverflow)?;
56        }
57        // Odd pitch is relatively prime to 32: rows at a fixed column map to
58        // distinct 32-bit shared-memory banks. Padding cells are never read.
59        let pitch = order + usize::from(order % 2 == 0);
60        let pivots = if kind == WarpDirectKind::Lu { order } else { 0 };
61        let shared_bytes = 4 * (order * pitch + order * rhs + pivots);
62        Ok(Self { batch, order, rhs, pitch, shared_bytes, kind })
63    }
64    /// Empty batches require no launch or plane capability, but shapes remain checked.
65    pub fn check_device(&self, limits: WarpDirectLimits) -> Result<(), WarpPlanError> {
66        if self.batch == 0 { return Ok(()); }
67        if limits.plane_min != 32 || limits.plane_max != 32 || !limits.plane_ops {
68            return Err(WarpPlanError::Unsupported("a fixed 32-lane plane with plane operations"));
69        }
70        if limits.max_threads < 32 || limits.max_block_x < 32 {
71            return Err(WarpPlanError::Unsupported("32 X threads per block"));
72        }
73        if self.batch > limits.max_grid_x as usize {
74            return Err(WarpPlanError::Unsupported("one X block per system; split this batch explicitly"));
75        }
76        if self.shared_bytes > limits.shared_bytes {
77            return Err(WarpPlanError::Unsupported("sufficient per-block shared memory"));
78        }
79        Ok(())
80    }
81    pub fn batch(&self) -> usize { self.batch }
82    pub fn order(&self) -> usize { self.order }
83    pub fn rhs(&self) -> usize { self.rhs }
84    pub fn pitch(&self) -> usize { self.pitch }
85    pub fn shared_bytes(&self) -> usize { self.shared_bytes }
86    pub fn kind(&self) -> WarpDirectKind { self.kind }
87    /// Successful-path logical global payload: A/B loaded once, factors/X stored
88    /// once, and status/pivots stored once. Excludes allocator, cache-line and
89    /// compiler-inserted traffic. Does NOT assert actual DRAM bytes transferred.
90    pub fn logical_global_bytes_per_system(&self) -> usize {
91        8 * (self.order * self.order + self.order * self.rhs) + 4
92            + if self.kind == WarpDirectKind::Lu { 4 * self.order } else { 0 }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    fn limits() -> WarpDirectLimits {
100        WarpDirectLimits { plane_min:32, plane_max:32, plane_ops:true, max_threads:1024,
101            max_block_x:1024, max_grid_x:65535, shared_bytes:48*1024 }
102    }
103    #[test] fn invalid_orders_and_rhs_are_rejected() {
104        for (n,r) in [(0,1),(33,1),(1,0),(1,9)] {
105            assert_eq!(WarpDirectPlan::new(1,n,r,WarpDirectKind::Lu),Err(WarpPlanError::InvalidShape));
106        }
107    }
108    #[test] fn overflow_is_rejected_without_allocating() {
109        assert!(WarpDirectPlan::new(usize::MAX,32,8,WarpDirectKind::Lu).is_err());
110        assert!(WarpDirectPlan::new(u32::MAX as usize/1024+1,32,8,WarpDirectKind::Lu).is_err());
111    }
112    #[test] fn shared_bounds_for_every_shape() {
113        for n in 1..=32 { for r in 1..=8 {
114            for kind in [WarpDirectKind::Cholesky,WarpDirectKind::Lu] {
115                let p=WarpDirectPlan::new(3,n,r,kind).unwrap();
116                assert!(p.pitch()>=n && p.pitch()<=n+1 && p.pitch()%2==1);
117                assert!(p.shared_bytes()<=5376);
118                assert_eq!(p.batch(),3); assert_eq!(p.kind(),kind);
119                p.check_device(limits()).unwrap();
120            }
121        }}
122    }
123    #[test] fn columns_have_no_bank_duplicates_in_this_layout_model() {
124        for n in 1..=32 {
125            let p=WarpDirectPlan::new(1,n,1,WarpDirectKind::Lu).unwrap();
126            for c in 0..n {
127                let mut banks=std::collections::BTreeSet::new();
128                for row in 0..n { assert!(banks.insert((row*p.pitch()+c)%32)); }
129            }
130        }
131    }
132    #[test] fn strided_copy_covers_only_logical_elements_once() {
133        for len in 1..=1024 {
134            let mut hits=vec![0;len];
135            for lane in 0..32 { for i in (lane..len).step_by(32) { hits[i]+=1; } }
136            assert!(hits.iter().all(|&x|x==1));
137        }
138    }
139    #[test] fn empty_batch_has_no_hardware_requirement() {
140        let p=WarpDirectPlan::new(0,32,8,WarpDirectKind::Lu).unwrap();
141        p.check_device(WarpDirectLimits{plane_min:0,plane_max:0,plane_ops:false,
142            max_threads:0,max_block_x:0,max_grid_x:0,shared_bytes:0}).unwrap();
143    }
144    #[test] fn each_hardware_requirement_is_checked() {
145        let p=WarpDirectPlan::new(65,32,8,WarpDirectKind::Lu).unwrap();
146        let mut bad=limits();bad.plane_min=16;assert!(p.check_device(bad).is_err());
147        bad=limits();bad.plane_max=64;assert!(p.check_device(bad).is_err());
148        bad=limits();bad.plane_ops=false;assert!(p.check_device(bad).is_err());
149        bad=limits();bad.max_threads=31;assert!(p.check_device(bad).is_err());
150        bad=limits();bad.max_block_x=31;assert!(p.check_device(bad).is_err());
151        bad=limits();bad.max_grid_x=64;assert!(p.check_device(bad).is_err());
152        bad=limits();bad.shared_bytes=5375;assert!(p.check_device(bad).is_err());
153        bad=limits();bad.shared_bytes=5376;p.check_device(bad).unwrap();
154    }
155    #[test] fn maximum_layout_and_payload_are_explicit() {
156        let a=WarpDirectPlan::new(1,32,8,WarpDirectKind::Cholesky).unwrap();
157        let b=WarpDirectPlan::new(1,32,8,WarpDirectKind::Lu).unwrap();
158        assert_eq!(a.shared_bytes(),5248);assert_eq!(b.shared_bytes(),5376);
159        assert_eq!(a.logical_global_bytes_per_system(),10244);
160        assert_eq!(b.logical_global_bytes_per_system(),10372);
161    }
162}
rusolver/kernel_plan.rs

rusolver/
kernel_plan.rs
