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

//! Reduction operations on GPU (sum, max, mean).

use super::device::GpuDevice;
use super::kernel::KernelCache;
use super::tensor::GpuTensor;

/// Reduce-sum a tensor along an axis.
pub fn reduce_sum(
    device: &GpuDevice,
    cache: &mut KernelCache,
    input: &GpuTensor,
    axis: usize,
) -> GpuTensor {
    reduce_op(device, cache, input, axis, ReduceOp::Sum)
}

/// Reduce-max a tensor along an axis.
pub fn reduce_max(
    device: &GpuDevice,
    cache: &mut KernelCache,
    input: &GpuTensor,
    axis: usize,
) -> GpuTensor {
    reduce_op(device, cache, input, axis, ReduceOp::Max)
}

/// Reduce-mean a tensor along an axis.
pub fn reduce_mean(
    device: &GpuDevice,
    cache: &mut KernelCache,
    input: &GpuTensor,
    axis: usize,
) -> GpuTensor {
    reduce_op(device, cache, input, axis, ReduceOp::Mean)
}

/// Sum all elements to a single scalar [1] tensor, entirely on GPU.
pub fn reduce_sum_all(
    device: &GpuDevice,
    cache: &mut KernelCache,
    input: &GpuTensor,
) -> GpuTensor {
    let mut t = reduce_sum(device, cache, input, 0);
    while t.numel() > 1 {
        t = reduce_sum(device, cache, &t, 0);
    }
    t
}

#[derive(Clone, Copy)]
enum ReduceOp {
    Sum,
    Max,
    Mean,
}

fn reduce_op(
    device: &GpuDevice,
    cache: &mut KernelCache,
    input: &GpuTensor,
    axis: usize,
    op: ReduceOp,
) -> GpuTensor {
    let shape = input.shape();
    assert!(axis < shape.len(), "reduce: axis out of bounds");

    let axis_size = shape[axis];
    let mut out_shape: Vec<usize> = shape.to_vec();
    out_shape.remove(axis);
    if out_shape.is_empty() {
        out_shape.push(1);
    }
    let out_numel: usize = out_shape.iter().product();

    // inner_size = product of dimensions after the axis
    let inner_size: usize = shape[axis + 1..].iter().product();

    let out = GpuTensor::uninit(device, &out_shape);

    // op_code: 0 = sum, 1 = max, 2 = mean
    let op_code = match op {
        ReduceOp::Sum => 0u32,
        ReduceOp::Max => 1,
        ReduceOp::Mean => 2,
    };

    let wgsl = r#"// Reduce along an axis: one thread per output element

struct Params {
    count: u32,
    axis_size: u32,
    inner_size: u32,
    op_code: u32,
}

@group(0) @binding(0) var<storage, read> input: array<f32>;
@group(0) @binding(1) var<storage, read_write> output: array<f32>;
@group(0) @binding(2) var<uniform> params: Params;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx >= params.count) { return; }

    let outer_idx = idx / params.inner_size;
    let inner_idx = idx % params.inner_size;
    let base = outer_idx * (params.axis_size * params.inner_size) + inner_idx;
    let stride = params.inner_size;

    var acc: f32;
    if (params.op_code == 1u) {
        acc = -3.402823e+38; // f32 MIN
    } else {
        acc = 0.0;
    }

    for (var k: u32 = 0u; k < params.axis_size; k = k + 1u) {
        let val = input[base + k * stride];
        if (params.op_code == 1u) {
            acc = max(acc, val);
        } else {
            acc = acc + val;
        }
    }

    if (params.op_code == 2u) {
        acc = acc / f32(params.axis_size);
    }

    output[idx] = acc;
}
"#;

    cache.dispatch_with_params(
        device,
        wgsl,
        &input.buffer,
        &out.buffer,
        &[out_numel as u32, axis_size as u32, inner_size as u32, op_code],
    );

    out
}