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.6 KiB116 linesraw
1//Tests automatic synthesis of Differential type requirement.
2
3//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type -dx12
4//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
5//TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute -shaderobj -output-using-type
6
7//TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer
8RWStructuredBuffer<float> outputBuffer;
9
10struct PathData : IDifferentiable
11{
12    float3 thp;    
13    uint length;
14    bool terminated;
15    bool isHit;
16
17    [BackwardDifferentiable]
18    __init()
19    {
20        this.thp = float3(1.f);
21        this.length = 0;
22        this.terminated = false;
23        this.isHit = false;
24    }
25}
26
27bool traceRayInline(uint length)
28{
29    if (length < 2) return true;
30    else return false;
31}
32
33float3 getAlbedo(uint length)
34{
35    return float3(0.9f, 1.f, 1.f);
36}
37
38float3 getAlbedoDerivative(uint length)
39{
40    return float3(1.f, 0.f, 0.f);
41}
42
43[ForwardDerivativeOf(getAlbedo)]
44[TreatAsDifferentiable]
45DifferentialPair<float3> __fwd_d_getAlbedo(uint length)
46{
47    float3 primalValue = getAlbedo(length);
48    float3 derivativeValue = no_diff getAlbedoDerivative(length);
49    return DifferentialPair<float3>(primalValue, derivativeValue);
50}
51
52[BackwardDerivativeOf(getAlbedo)]
53[TreatAsDifferentiable]
54void __bwd_d_getAlbedo(uint length, float3 dOut)
55{
56    outputBuffer[2] += dOut.x;
57}
58
59[BackwardDifferentiable]
60void handleHit(inout PathData pathData)
61{
62    if (pathData.length >= 2)
63    {
64        pathData.terminated = true;
65        return;
66    }
67
68    float3 albedo = getAlbedo(pathData.length);
69    pathData.thp *= albedo;
70    pathData.length++;
71}
72
73[BackwardDifferentiable]
74[PreferRecompute]
75float3 tracePath()
76{
77    PathData pathData = PathData();
78
79    if (traceRayInline(pathData.length))
80    {
81        pathData.isHit = true;
82    }
83    else
84    {
85        pathData.terminated = true;
86        pathData.isHit = false;
87    }
88
89    [MaxIters(4)]
90    while (!pathData.terminated)
91    {
92        if (pathData.isHit)
93        {
94            handleHit(pathData);
95
96            //pathData.isHit = traceRayInline(pathData.length);
97            if (!traceRayInline(pathData.length)) pathData.isHit = false;
98            else pathData.isHit = true;
99        }
100        else
101        {
102            pathData.terminated = true;
103        }
104    }
105    return pathData.thp;
106}
107
108[numthreads(1, 1, 1)]
109void computeMain(uint3 dispathThreadID: SV_DispatchThreadID)
110{
111    DifferentialPair<float3> dpThp = __fwd_diff(tracePath)();
112    outputBuffer[0] = dpThp.p.x;
113    outputBuffer[1] = dpThp.d.x;
114
115    __bwd_diff(tracePath)(float3(1.f, 0.f, 0.f));
116}