spark_runtime/
gpu_args.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Packing typed kernel arguments into the driver's u64 parameter slots.
4//!
5//! Split out of `gpu.rs` for the repo's 500-line cap. It earns its own file: the
6//! packing is the one place that knows an argument may span MORE than one slot,
7//! and the version this replaced silently truncated anything wider than 8 bytes.
8
9use crate::gpu::KernelArg;
10
11/// Pack typed args into u64 slots, returning the slots and each argument's
12/// STARTING slot index.
13///
14/// Slots are u64-granular, but an argument is NOT limited to one slot: a
15/// by-value struct parameter (`CUtensorMap` is 128 bytes) occupies
16/// `ceil(len/8)` CONSECUTIVE slots and contributes exactly ONE entry to the
17/// param array, pointing at the first.
18///
19/// ★ The packing this replaces was `let n = b.len().min(8)` — anything wider
20/// was TRUNCATED TO ITS FIRST 8 BYTES AND LAUNCHED, with no error. Nothing
21/// passed more than 8 bytes yet, so it never fired; the first caller that did
22/// would have got a kernel reading garbage out of a struct parameter and no
23/// diagnostic anywhere. Silent truncation is not an acceptable failure mode on
24/// the launch path, so this is a free function with its own tests rather than
25/// eight lines buried in a default trait method.
26pub fn pack_kernel_args(args: &[KernelArg<'_>]) -> (Vec<u64>, Vec<usize>) {
27    let total_slots: usize = args
28        .iter()
29        .map(|a| match a {
30            KernelArg::Buffer(_) => 1,
31            KernelArg::Bytes(b) => b.len().div_ceil(8).max(1),
32        })
33        .sum();
34    let mut storage: Vec<u64> = Vec::with_capacity(total_slots);
35    let mut starts: Vec<usize> = Vec::with_capacity(args.len());
36    for arg in args {
37        starts.push(storage.len());
38        match arg {
39            KernelArg::Buffer(p) => storage.push(p.0),
40            KernelArg::Bytes(b) => {
41                for c in b.chunks(8) {
42                    let mut slot = [0u8; 8];
43                    slot[..c.len()].copy_from_slice(c);
44                    storage.push(u64::from_le_bytes(slot));
45                }
46                if b.is_empty() {
47                    storage.push(0);
48                }
49            }
50        }
51    }
52    (storage, starts)
53}
spark_runtime/gpu_args.rs

spark_runtime/
gpu_args.rs
