yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Yong HeSmall fix to buffer load specialization pass to allow more specialization to happen. (#8653)3cf1f5a61

master
branchesmasterall refs →
1.3 KiB56 linesraw
1//TEST:SIMPLE(filecheck=CHECK):-target spirv
2
3// CHECK: OpEntryPoint
4
5// Make sure we never load the entire TensorList struct to local registers,
6// instead, we should specialize the fetch function to directly load the
7// element tensor from gTensors.
8
9// CHECK-NOT: OpLoad %TensorList_std140
10
11struct RWTensor<T, let D : int>
12{
13    int dims[D];
14    RWStructuredBuffer<T> buffer;
15    T getv(vector<uint, D> index)
16    {
17        int flat_index = 0;
18        int stride = 1;
19        for (int i = D - 1; i >= 0; --i)
20        {
21            flat_index += index[i] * stride;
22            stride *= dims[i];
23        }
24        return buffer[flat_index];
25    }
26}
27struct TensorList<let N : int>
28{
29    RWTensor<float, 2> tensors[N];
30
31    float fetch(int tensor_index, uint2 index)
32    {
33        return tensors[tensor_index].getv(index);
34    }
35}
36
37float sum_indirect<let N : int>(uint2 tid, TensorList<N> tensor_list, uint tensor_indices[N])
38{
39    float result = 0.0;
40    for (int i = 0; i < N; i++)
41    {
42        result += tensor_list.fetch(tensor_indices[i], tid);
43    }
44    return result;
45}
46
47uniform TensorList<32> gTensors;
48uniform uint gTensorIndices[32];
49
50uniform float* result;
51
52[numthreads(1,1,1)]
53void computeMain(uint2 tid : SV_DispatchThreadID)
54{
55    *result = sum_indirect(tid, gTensors, gTensorIndices);
56}