tang 0.2.0

Math library for physical reality — geometry, spatial algebra, tensor, training, GPU compute, and 3D gaussian splatting
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! GPU buffer: upload, download, staging.

use super::device::GpuDevice;

/// A GPU storage buffer holding f32 values.
pub struct GpuBuffer {
    pub(crate) buffer: wgpu::Buffer,
    pub(crate) len: usize,
}

impl GpuBuffer {
    /// Create a storage buffer initialized from a slice.
    pub fn from_slice(device: &GpuDevice, data: &[f32]) -> Self {
        use wgpu::util::DeviceExt;
        let buffer = device
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("tang-gpu storage"),
                contents: bytemuck::cast_slice(data),
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
            });
        Self {
            buffer,
            len: data.len(),
        }
    }

    /// Create a storage buffer initialized from a u32 slice.
    ///
    /// Used for token ID buffers passed to [`GpuEmbedding`].
    /// The buffer holds u32 values but `len` tracks the element count.
    pub fn from_u32_slice(device: &GpuDevice, data: &[u32]) -> Self {
        use wgpu::util::DeviceExt;
        let buffer = device
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("tang-gpu storage u32"),
                contents: bytemuck::cast_slice(data),
                usage: wgpu::BufferUsages::STORAGE
                    | wgpu::BufferUsages::COPY_SRC
                    | wgpu::BufferUsages::COPY_DST,
            });
        Self {
            buffer,
            len: data.len(),
        }
    }

    /// Create an uninitialized storage buffer of `len` f32 elements.
    pub fn uninit(device: &GpuDevice, len: usize) -> Self {
        let buffer = device.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tang-gpu storage"),
            size: (len * std::mem::size_of::<f32>()) as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        Self { buffer, len }
    }

    /// Download buffer contents to CPU.
    pub async fn to_vec(&self, device: &GpuDevice) -> Vec<f32> {
        let size = (self.len * std::mem::size_of::<f32>()) as u64;
        let staging = device.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tang-gpu staging"),
            size,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let mut encoder = device
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("tang-gpu download"),
            });
        encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging, 0, size);
        device.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).ok();
        });
        device.device.poll(wgpu::Maintain::Wait);
        rx.recv().unwrap().unwrap();

        let data = slice.get_mapped_range();
        let result: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
        drop(data);
        staging.unmap();
        result
    }

    /// Download buffer contents synchronously.
    pub fn to_vec_sync(&self, device: &GpuDevice) -> Vec<f32> {
        pollster::block_on(self.to_vec(device))
    }

    /// Number of f32 elements.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Whether the buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Size in bytes.
    pub fn byte_size(&self) -> u64 {
        (self.len * std::mem::size_of::<f32>()) as u64
    }

    /// Clone this buffer entirely on the GPU (no CPU readback).
    /// Submits immediately — prefer `clone_gpu_batched` when batching commands.
    pub fn clone_gpu(&self, device: &GpuDevice) -> GpuBuffer {
        let dst = Self::uninit(device, self.len);
        let mut encoder = device
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("tang-gpu clone"),
            });
        encoder.copy_buffer_to_buffer(&self.buffer, 0, &dst.buffer, 0, self.byte_size());
        device.queue.submit(std::iter::once(encoder.finish()));
        dst
    }

    /// Clone this buffer on GPU, using the kernel cache's batching mode.
    pub fn clone_gpu_batched(&self, device: &GpuDevice, cache: &mut super::kernel::KernelCache) -> GpuBuffer {
        let dst = Self::uninit(device, self.len);
        let mut encoder = device
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("tang-gpu clone"),
            });
        encoder.copy_buffer_to_buffer(&self.buffer, 0, &dst.buffer, 0, self.byte_size());
        cache.submit_or_enqueue(device, encoder.finish());
        dst
    }
}