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 KiB60 linesraw
1
2//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK): -slang -compute -shaderobj -output-using-type
3//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-cuda -compute -shaderobj -output-using-type
4
5//TEST_INPUT:ubuffer(data=[0 0 0 0 0], stride=4):out,name=outputBuffer
6RWStructuredBuffer<float> outputBuffer;
7
8
9[BackwardDerivative(set_bwd)]
10void set(uint idx, float x)
11{
12    outputBuffer[idx] = x;
13}
14
15void set_bwd(uint idx, inout DifferentialPair<float> x)
16{
17    // For debugging, we'll set the derivative to 1.0
18    x = DifferentialPair<float>(x.p, 1.0f);
19}
20
21[Differentiable]
22void run(
23    uint idx,
24    float x)
25{
26    if (idx >= 1) return;
27
28    if (idx == 0)
29    { }
30
31    for (int i = 0; i < 1; i++)
32    {
33        if (idx > 0)
34        {
35            return;
36        }
37
38        if (idx == 0)
39        { 
40            x = x * 2.0f;
41        }
42    }
43
44    if (idx == 0)
45    { }
46
47    set(idx, x);
48}
49
50[numthreads(1, 1, 1)]
51void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
52{
53    // bwd_diff
54    DifferentialPair<float> dpa = DifferentialPair<float>(1.0, 0.0);
55    bwd_diff(run)(dispatchThreadID.x, dpa);
56    outputBuffer[dispatchThreadID.x] = dpa.d; 
57    
58    // CHECK: type: float
59    // CHECK: 2.0
60}