yum-mirror/slang

Making it easier to work with shaders

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

Sai Praveen BangaruAdd checking for differentiability of the primal substitute function. (#6277)551bbb5fb

master
branchesmasterall refs →
1.1 KiB44 linesraw
1//DIAGNOSTIC_TEST:SIMPLE(filecheck=CHECK): -target spirv -entry computeMain -stage compute
2
3//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBufferPrimal
4RWStructuredBuffer<float> outputBufferPrimal;
5
6//TEST_INPUT:ubuffer(data=[1 2 3 4], stride=4):name=gradBuffer
7RWStructuredBuffer<float> gradBuffer;
8
9struct BufferWithGrad
10{
11    RWStructuredBuffer<float> primal;
12    RWStructuredBuffer<float> grad;
13
14    [Differentiable]
15    void add(float value) { primal[0] = primal[0] + detach(value); }
16
17    // check for diagnostic:
18    // CHECK-DAG: ([[# @LINE+1]]): error 31158
19    [PrimalSubstituteOf(add)]
20    void add_subst(float value)
21    {
22    }
23
24    [BackwardDerivativeOf(add)]
25    void add_bwd(inout DifferentialPair<float> d)
26    {
27        d = diffPair(d.p, grad[0]);
28    }
29}
30
31[Differentiable]
32void diffCall(BufferWithGrad result)
33{
34    result.add(1.0f);
35}
36
37[shader("compute")]
38[numthreads(1, 1, 1)]
39void computeMain(uint3 dispatchThreadID: SV_DispatchThreadID)
40{
41    BufferWithGrad bg = {outputBufferPrimal, gradBuffer};
42    diffCall(bg); 
43    bwd_diff(diffCall)(bg);
44}