yum-mirror/slang

Making it easier to work with shaders

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

Harsh Aggarwal (NVIDIA)Fix 7723 - Add autodiff tests (#7919)d10732742

master
branchesmasterall refs →
1.2 KiB42 linesraw
1//TEST:COMPARE_COMPUTE(filecheck-buffer=CHECK):-output-using-type
2//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-cuda -compute -shaderobj -output-using-type
3
4struct GradientBuffer
5{
6    StructuredBuffer<float> grads;
7
8    [Differentiable]
9    void write(int idx, float value) { /* Discard write */ }
10
11    [BackwardDerivativeOf(write)]
12    void write_bwd(int idx, inout DifferentialPair<float> d)
13    {
14        d = diffPair(d.p, grads[idx]);
15    }
16}
17
18[Differentiable]
19void test(GradientBuffer dst, int idx, float v)
20{
21    dst.write(idx, v);
22}
23
24//TEST_INPUT: set grad_in = ubuffer(data=[101.0 102.0 103.0 104.0], stride=4)
25uniform StructuredBuffer<float> grad_in;
26
27//TEST_INPUT: set grad_out = ubuffer(data=[0 0 0 0], stride=4)
28uniform RWStructuredBuffer<float> grad_out;
29
30//TEST_INPUT: set output = out ubuffer(data=[0 0 0 0], stride=4)
31uniform RWStructuredBuffer<float> output;
32
33[shader("compute")]
34[numthreads(1,1,1)]
35void computeMain()
36{
37    GradientBuffer grads = { grad_in };
38    DifferentialPair<float> result = diffPair(1.0f);
39    bwd_diff(test)(grads, 0, result);
40    // CHECK: 101.0
41    output[0] = result.d; // Should return grad_in[0], but returns 0.0f instead
42}