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 →
4.9 KiB225 linesraw
1//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type
2//TEST(compute):COMPARE_COMPUTE_EX:-slang -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 0], stride=4):out,name=outputBuffer
6RWStructuredBuffer<float> outputBuffer;
7
8typedef float Real;
9
10__generic<let N : int>
11struct myvector
12{
13    vector<Real, N> val;
14 
15    [TreatAsDifferentiable]
16    __init(vector<Real,N> data)
17    {
18        val = data;
19    }
20}
21
22extension myvector<3> : MyLinearArithmeticType
23{
24    [ForwardDifferentiable]
25    static myvector<3> ladd(myvector<3> a, myvector<3> b)
26    {
27        return myvector<3>(a.val + b.val);
28    }
29
30    [ForwardDifferentiable]
31    static myvector<3> lmul(myvector<3> a, myvector<3> b)
32    {
33        return myvector<3>(a.val * b.val);
34    }
35
36    [ForwardDifferentiable]
37    static myvector<3> lscale(float a, myvector<3> b)
38    {
39        return myvector<3>(a * b.val);
40    }
41
42    [ForwardDifferentiable]
43    static float ldot(myvector<3> a, myvector<3> b)
44    {
45        return dot(a.val, b.val);
46    }
47
48    [ForwardDifferentiable]
49    __init(vector<Real, 3> a)
50    {
51        val = a;
52    }
53};
54
55
56extension myvector<4> : MyLinearArithmeticType
57{
58    [ForwardDifferentiable]
59    static myvector<4> ladd(myvector<4> a, myvector<4> b)
60    {
61        return myvector<4>(a.val + b.val);
62    }
63
64    [ForwardDifferentiable]
65    static myvector<4> lmul(myvector<4> a, myvector<4> b)
66    {
67        return myvector<4>(a.val * b.val);
68    }
69
70    [ForwardDifferentiable]
71    static myvector<4> lscale(float a, myvector<4> b)
72    {
73        return myvector<4>(a * b.val);
74    }
75
76    [ForwardDifferentiable]
77    static float ldot(myvector<4> a, myvector<4> b)
78    {
79        return dot(a.val, b.val);
80    }
81
82    [ForwardDifferentiable]
83    __init(vector<Real, 4> a)
84    {
85        val = a;
86    }
87
88};
89
90typedef myvector<3> myfloat3;
91typedef myvector<4> myfloat4;
92
93typedef DifferentialPair<Real> dpfloat;
94
95[TreatAsDifferentiable]
96interface MyLinearArithmeticType
97{
98    static This ladd(This a, This b);
99    static This lmul(This a, This b);
100    static This lscale(Real a, This b);
101    static Real ldot(This a, This b);
102};
103
104extension myfloat3 : IDifferentiable
105{
106    typedef myfloat3 Differential;
107
108    [DerivativeMember(Differential.val)]
109    extern vector<Real, 3> val;
110
111    static Differential dzero()
112    {
113        return myfloat3(0);
114    }
115
116    [ForwardDifferentiable]
117    static Differential dadd(Differential a, Differential b)
118    {
119        return a + b;
120    }
121
122    [ForwardDifferentiable]
123    static Differential dmul<T : __BuiltinRealType>(T a, Differential b)
124    {
125        return myfloat3(__realCast<Real, T>(a) * b.val);
126    }
127
128};
129
130extension myfloat4 : IDifferentiable
131{
132    typedef myfloat4 Differential;
133
134    [DerivativeMember(Differential.val)]
135    extern vector<Real, 4> val;
136
137    static Differential dzero()
138    {
139        return myfloat4(0);
140    }
141
142    [ForwardDifferentiable]
143    static Differential dadd(Differential a, Differential b)
144    {
145        return a + b;
146    }
147
148    [ForwardDifferentiable]
149    static Differential dmul<T: __BuiltinRealType>(T a, Differential b)
150    {
151        return myfloat4(__realCast<Real, T>(a) * b.val);
152    }
153};
154
155typedef DifferentialPair<myfloat4> dpfloat4;
156typedef DifferentialPair<myfloat3> dpfloat3;
157
158extension float : MyLinearArithmeticType
159{
160    [ForwardDifferentiable]
161    static float ladd(float a, float b)
162    {
163        return a + b;
164    }
165
166    [ForwardDifferentiable]
167    static float lmul(float a, float b)
168    {
169        return a * b;
170    }
171
172    [ForwardDifferentiable]
173    static float lscale(float a, float b)
174    {
175        return a * b;
176    }
177    
178    [ForwardDifferentiable]
179    static float ldot(float a, float b)
180    {
181        return a * b;
182    }
183};
184
185typealias MyLinearArithmeticDifferentiableType = IDifferentiable & MyLinearArithmeticType;
186
187__generic<T : MyLinearArithmeticDifferentiableType>
188[ForwardDifferentiable]
189T operator +(T a, T b)
190{
191    return T.ladd(a, b);
192}
193
194__generic<T : MyLinearArithmeticDifferentiableType>
195[ForwardDifferentiable]
196T operator *(T a, T b)
197{
198    return T.lmul(a, b);
199}
200
201__generic<G : MyLinearArithmeticDifferentiableType>
202[ForwardDifferentiable]
203G f(G x)
204{
205    G a = x + x;
206    G b = x * x;
207
208    return a * a + G.lscale((Real)3.0, x); 
209}
210
211
212[numthreads(1, 1, 1)]
213void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
214{
215    {
216        dpfloat dpa = dpfloat(2.0, 1.0);
217        dpfloat4 dpf4 = dpfloat4(myfloat4(float4(1.5, 2.0, 0.5, 1.0)), myfloat4(float4(0.5, 0.8, 1.6, 2.5)));
218        dpfloat3 dpf3 = dpfloat3(myfloat3(float3(1.0, 3.0, 5.0)), myfloat3(float3(0.5, 1.5, 2.5)));
219
220        outputBuffer[0] = f(dpa.p);                       // Expect: 22.0
221        outputBuffer[1] = __fwd_diff(f)(dpfloat(2.0, 0.5)).d;  // Expect: 9.5
222        outputBuffer[2] = __fwd_diff(f)(dpf4).d.val.w;         // Expect: 27.5
223        outputBuffer[3] = __fwd_diff(f)(dpf3).d.val.y;         // Expect: 40.5
224    }
225}