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 →
2.2 KiB78 linesraw
1//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type
2//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
3//TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute -shaderobj -output-using-type
4
5//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
6RWStructuredBuffer<float> outputBuffer;
7
8//TEST_INPUT:ubuffer(data=[0.0 1.0 2.0 3.0], stride=4):name=endpointBuffer
9RWStructuredBuffer<float> endpointBuffer;
10
11//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=endpointDifferentialBuffer
12RWStructuredBuffer<float> endpointDifferentialBuffer;
13
14struct LineSegment : IDifferentiable
15{
16    float x0;
17    float x1;
18
19    [BackwardDifferentiable]
20    __init(float _x0, float _x1)
21    {
22        x0 = _x0;
23        x1 = _x1;
24    }
25};
26
27[BackwardDerivative(d_loadLineSegment)]
28[ForwardDerivative(fwd_loadLineSegment)]
29LineSegment loadLineSegment(uint id)
30{
31    return {endpointBuffer[id * 2], endpointBuffer[id * 2 + 1]};
32}
33
34[BackwardDerivative(d_fwd_loadLineSegment)]
35DifferentialPair<LineSegment> fwd_loadLineSegment(uint id)
36{
37    return DifferentialPair<LineSegment>(loadLineSegment(id), LineSegment.dzero());
38}
39
40void d_loadLineSegment(uint id, LineSegment.Differential d_ls)
41{
42    endpointDifferentialBuffer[id * 2] += d_ls.x0;
43    endpointDifferentialBuffer[id * 2 + 1] += d_ls.x1;
44}
45
46void d_fwd_loadLineSegment(uint id, DifferentialPair<LineSegment>.Differential dp_ls)
47{
48    endpointDifferentialBuffer[id * 2] += dp_ls.p.x0;
49    endpointDifferentialBuffer[id * 2 + 1] += dp_ls.p.x1;
50}
51
52[BackwardDifferentiable]
53float something()
54{
55    LineSegment ls = __fwd_diff(loadLineSegment)(1).p;
56    return ls.x0 + ls.x1;
57}
58
59[numthreads(1, 1, 1)]
60void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
61{
62    {
63        LineSegment ls = __fwd_diff(loadLineSegment)(0).p;
64        outputBuffer[0] = ls.x0; // Expect: 0
65        outputBuffer[1] = ls.x1; // Expect: 1
66    }
67
68    {
69        LineSegment.Differential d_ls = __fwd_diff(loadLineSegment)(0).d;
70        outputBuffer[2] = d_ls.x1; // Expect: 0
71    }
72
73    {
74        // Expect: 2.0 in endpointDifferentialBuffer[2]
75        // Expect: 2.0 in endpointDifferentialBuffer[3]
76        __bwd_diff(something)(2.0); 
77    }
78}