runmat-accelerate 0.4.1

Pluggable GPU acceleration layer for RunMat (CUDA, ROCm, Metal, Vulkan/Spir-V)
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

pub const TRIL_SHADER_F64: &str = r#"
struct Tensor {
    data: array<f64>,
};

struct TrilParams {
    len: u32,
    start: u32,
    rows: u32,
    cols: u32,
    plane: u32,
    diag_offset: i32,
    _pad0: u32,
    _pad1: u32,
};

@group(0) @binding(0) var<storage, read> input0: Tensor;
@group(0) @binding(1) var<storage, read_write> output: Tensor;
@group(0) @binding(2) var<uniform> params: TrilParams;

@compute @workgroup_size(@WG@)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    if gid.x >= params.len {
        return;
    }
    let index: u32 = params.start + gid.x;
    let plane = params.plane;
    if plane == 0u {
        output.data[index] = input0.data[index];
        return;
    }
    let rows = params.rows;
    let within = index % plane;
    let row = within % rows;
    let col = within / rows;

    let row_i = i32(row);
    let col_i = i32(col);
    var zero_out = false;
    if params.diag_offset >= 0 {
        let threshold = col_i - params.diag_offset;
        zero_out = row_i < threshold;
    } else {
        let k = -params.diag_offset;
        let diff = row_i - col_i;
        zero_out = diff < k;
    }

    if zero_out {
        output.data[index] = 0.0;
    } else {
        output.data[index] = input0.data[index];
    }
}
"#;

pub const TRIL_SHADER_F32: &str = r#"
struct Tensor {
    data: array<f32>,
};

struct TrilParams {
    len: u32,
    start: u32,
    rows: u32,
    cols: u32,
    plane: u32,
    diag_offset: i32,
    _pad0: u32,
    _pad1: u32,
};

@group(0) @binding(0) var<storage, read> input0: Tensor;
@group(0) @binding(1) var<storage, read_write> output: Tensor;
@group(0) @binding(2) var<uniform> params: TrilParams;

@compute @workgroup_size(@WG@)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    if gid.x >= params.len {
        return;
    }
    let index: u32 = params.start + gid.x;
    let plane = params.plane;
    if plane == 0u {
        output.data[index] = input0.data[index];
        return;
    }
    let rows = params.rows;
    let within = index % plane;
    let row = within % rows;
    let col = within / rows;

    let row_i = i32(row);
    let col_i = i32(col);
    var zero_out = false;
    if params.diag_offset >= 0 {
        let threshold = col_i - params.diag_offset;
        zero_out = row_i < threshold;
    } else {
        let k = -params.diag_offset;
        let diff = row_i - col_i;
        zero_out = diff < k;
    }

    if zero_out {
        output.data[index] = 0.0;
    } else {
        output.data[index] = input0.data[index];
    }
}
"#;