silvic_gpu/
params.rs

1//! The kernels' parameter structs, mirrored field for field from `kernels/common.h`, and the capacities both sides
2//! share. Every field is 4 bytes, so `#[repr(C)]` lays the structs out exactly as Metal and CUDA do.
3
4use bytemuck::{Pod, Zeroable};
5use silvic_core::data::GROUP_BIN_BUDGET;
6
7/// Bins of one quantized histogram work unit: a whole feature group, one packed i32 per bin (32 KiB of threadgroup
8/// memory, Metal's limit).
9pub(crate) const HIST_QUANT_BINS: u32 = GROUP_BIN_BUDGET;
10/// Bins of one float histogram work unit: separate f32 g and h arrays fill the same 32 KiB.
11pub(crate) const HIST_FLOAT_BINS: u32 = GROUP_BIN_BUDGET / 2;
12/// Threads of the reduction kernels (`grad_stats`, `leaf_sums`); a power of two.
13pub(crate) const REDUCE_THREADS: u32 = 256;
14/// Threads of the compaction kernels (`bag_*`, `partition_*`, `scan_blocks`); a power of two.
15pub(crate) const SCAN_THREADS: u32 = 256;
16/// Largest query each lambdarank kernel holds in threadgroup memory, and the threads it runs with.
17pub(crate) const LAMBDARANK_CLASSES: [(u32, u32); 3] = [(64, 64), (256, 256), (1024, 256)];
18/// Threads of the widest lambdarank kernel: the size of its reduction array.
19pub(crate) const LAMBDARANK_THREADS_MAX: u32 = 256;
20/// `HistParams::side`: the leaf's own rows, or its left or right child after a partition encoded just before.
21pub(crate) const HIST_SIDE_LEAF: u32 = 0;
22pub(crate) const HIST_SIDE_LEFT: u32 = 1;
23pub(crate) const HIST_SIDE_RIGHT: u32 = 2;
24
25const COMMON: &str = include_str!("../kernels/common.h");
26const KERNELS: &str = include_str!("../kernels/kernels.kern");
27
28/// The complete kernel source for one API: `prelude`, the shared capacities as `#define`s, `common.h` and
29/// `kernels.kern`.
30pub fn kernel_source(prelude: &str) -> String {
31    let defines = [
32        ("HIST_QUANT_BINS", HIST_QUANT_BINS),
33        ("HIST_FLOAT_BINS", HIST_FLOAT_BINS),
34        ("REDUCE_THREADS", REDUCE_THREADS),
35        ("SCAN_THREADS", SCAN_THREADS),
36        ("LAMBDARANK_CAP_SMALL", LAMBDARANK_CLASSES[0].0),
37        ("LAMBDARANK_CAP_MEDIUM", LAMBDARANK_CLASSES[1].0),
38        ("LAMBDARANK_CAP_LARGE", LAMBDARANK_CLASSES[2].0),
39        ("LAMBDARANK_THREADS_MAX", LAMBDARANK_THREADS_MAX),
40        ("HIST_SIDE_LEFT", HIST_SIDE_LEFT),
41        ("HIST_SIDE_RIGHT", HIST_SIDE_RIGHT),
42    ];
43    let mut source = String::with_capacity(prelude.len() + COMMON.len() + KERNELS.len() + 512);
44    source.push_str(prelude);
45    source.push('\n');
46    for (name, value) in defines {
47        source.push_str(&format!("#define {name} {value}u\n"));
48    }
49    source.push_str(COMMON);
50    source.push('\n');
51    source.push_str(KERNELS);
52    source
53}
54
55#[repr(C)]
56#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
57pub(crate) struct FillParams {
58    pub n: u32,
59    pub value: u32,
60}
61
62#[repr(C)]
63#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
64pub(crate) struct GradParams {
65    pub n: u32,
66    pub has_weight: u32,
67    pub sigmoid: f32,
68    pub pad: u32,
69}
70
71#[repr(C)]
72#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
73pub(crate) struct StatsParams {
74    pub n: u32,
75    pub pad: [u32; 3],
76}
77
78#[repr(C)]
79#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
80pub(crate) struct QuantParams {
81    pub n: u32,
82    pub row_offset_lo: u32,
83    pub row_offset_hi: u32,
84    pub inv_grad_scale: f32,
85    pub inv_hess_scale: f32,
86    pub grad_max_q: i32,
87    pub hess_max_q: i32,
88    pub constant_hessian: u32,
89    pub stochastic: u32,
90    pub seed: u32,
91    pub iteration: u32,
92    pub pad: u32,
93}
94
95#[repr(C)]
96#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
97pub(crate) struct BagParams {
98    pub n: u32,
99    pub chunk: u32,
100    pub row_offset_lo: u32,
101    pub row_offset_hi: u32,
102    pub seed: u32,
103    pub round: u32,
104    pub threshold: u32,
105    pub pad: u32,
106}
107
108#[repr(C)]
109#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
110pub(crate) struct ScanParams {
111    pub n: u32,
112    pub pad: [u32; 3],
113}
114
115#[repr(C)]
116#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
117pub(crate) struct PartitionParams {
118    pub begin: u32,
119    pub count: u32,
120    pub chunk: u32,
121    pub row_stride: u32,
122    pub group_base_lo: u32,
123    pub group_base_hi: u32,
124    pub column: u32,
125    pub pad: u32,
126    pub left_bins: [u32; 8],
127}
128
129#[repr(C)]
130#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
131pub(crate) struct CopyParams {
132    pub begin: u32,
133    pub count: u32,
134    pub pad: [u32; 2],
135}
136
137#[repr(C)]
138#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
139pub(crate) struct HistUnit {
140    pub words_base_lo: u32,
141    pub words_base_hi: u32,
142    pub row_words: u32,
143    pub bin_begin: u32,
144    pub num_bins: u32,
145    pub cols_begin: u32,
146    pub pad: [u32; 2],
147}
148
149#[repr(C)]
150#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
151pub(crate) struct HistParams {
152    pub begin: u32,
153    pub count: u32,
154    pub rows_per_block: u32,
155    pub side: u32,
156}
157
158#[repr(C)]
159#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
160pub(crate) struct LeafSumParams {
161    pub grad_scale: f32,
162    pub hess_scale: f32,
163    pub pad: [u32; 2],
164}
165
166#[repr(C)]
167#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
168pub(crate) struct GroupLayout {
169    pub base_lo: u32,
170    pub base_hi: u32,
171    pub row_stride: u32,
172    pub pad: u32,
173}
174
175#[repr(C)]
176#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
177pub(crate) struct TreeNode {
178    pub group: u32,
179    pub column: u32,
180    pub left: i32,
181    pub right: i32,
182    pub left_bins: [u32; 8],
183}
184
185#[repr(C)]
186#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
187pub(crate) struct TreeParams {
188    pub n: u32,
189    pub num_nodes: u32,
190    pub pad: [u32; 2],
191}
192
193#[repr(C)]
194#[derive(Debug, Clone, Copy, Default, Pod, Zeroable)]
195pub(crate) struct RankParams {
196    pub query_begin: u32,
197    pub num_queries: u32,
198    pub has_weight: u32,
199    pub norm: u32,
200    pub truncation_level: u32,
201    pub sigmoid: f32,
202    pub sigmoid_min_hi: f32,
203    pub sigmoid_min_lo: f32,
204    pub sigmoid_factor_hi: f32,
205    pub sigmoid_factor_lo: f32,
206    pub sigmoid_last: u32,
207    pub pad: u32,
208}
209
210// The sizes of the matching structs in `kernels/common.h`.
211const _: () = assert!(std::mem::size_of::<FillParams>() == 8);
212const _: () = assert!(std::mem::size_of::<GradParams>() == 16);
213const _: () = assert!(std::mem::size_of::<StatsParams>() == 16);
214const _: () = assert!(std::mem::size_of::<QuantParams>() == 48);
215const _: () = assert!(std::mem::size_of::<BagParams>() == 32);
216const _: () = assert!(std::mem::size_of::<ScanParams>() == 16);
217const _: () = assert!(std::mem::size_of::<PartitionParams>() == 64);
218const _: () = assert!(std::mem::size_of::<CopyParams>() == 16);
219const _: () = assert!(std::mem::size_of::<HistUnit>() == 32);
220const _: () = assert!(std::mem::size_of::<HistParams>() == 16);
221const _: () = assert!(std::mem::size_of::<LeafSumParams>() == 16);
222const _: () = assert!(std::mem::size_of::<GroupLayout>() == 16);
223const _: () = assert!(std::mem::size_of::<TreeNode>() == 48);
224const _: () = assert!(std::mem::size_of::<TreeParams>() == 16);
225const _: () = assert!(std::mem::size_of::<RankParams>() == 48);
silvic_gpu/params.rs

silvic_gpu/
params.rs
