metrale_gpu_runtime/
gpu_args.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Packing typed kernel arguments into the driver's u64 parameter
4//! slots.
5//!
6//! Owner: gpu-runtime.
7//! Invariants:
8//! - `starts` has one entry per argument, in order, pointing at its first slot.
9//! - Every byte of every argument is copied; none is truncated.
10
11use crate::gpu::KernelArg;
12
13/// 2026-09-25: Pack typed args into u64 slots; returns the slots and each
14/// argument's starting slot index.
15///
16/// A buffer takes one slot. A byte argument takes `ceil(len / 8)` consecutive
17/// slots (one when empty), little-endian, the last zero-padded; a 128-byte
18/// `CUtensorMap` takes 16. Each argument contributes one entry to `starts`.
19pub fn pack_kernel_args(args: &[KernelArg<'_>]) -> (Vec<u64>, Vec<usize>) {
20    let total_slots: usize = args
21        .iter()
22        .map(|a| match a {
23            KernelArg::Buffer(_) => 1,
24            KernelArg::Bytes(b) => b.len().div_ceil(8).max(1),
25        })
26        .sum();
27    let mut storage: Vec<u64> = Vec::with_capacity(total_slots);
28    let mut starts: Vec<usize> = Vec::with_capacity(args.len());
29    for arg in args {
30        starts.push(storage.len());
31        match arg {
32            KernelArg::Buffer(p) => storage.push(p.0),
33            KernelArg::Bytes(b) => {
34                for c in b.chunks(8) {
35                    let mut slot = [0u8; 8];
36                    slot[..c.len()].copy_from_slice(c);
37                    storage.push(u64::from_le_bytes(slot));
38                }
39                if b.is_empty() {
40                    storage.push(0);
41                }
42            }
43        }
44    }
45    (storage, starts)
46}
metrale_gpu_runtime/gpu_args.rs

metrale_gpu_runtime/
gpu_args.rs
