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 →
8.1 KiB342 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
10typealias IDFloat = __BuiltinRealType & IDifferentiable;
11
12__generic<T : IDifferentiable, let N : int>
13struct dvector : IDifferentiable
14{
15    typedef dvector<T.Differential, N> Differential;
16    [DerivativeMember(Differential.values)]
17    T values[N];
18};
19
20__generic<T : IDFloat, let N : int>
21struct myvector : IDifferentiable
22{
23    typedef dvector<T.Differential, N> Differential;
24
25    [DerivativeMember(Differential.values)]
26    T values[N];
27    
28    __init(T c)
29    {
30        [ForceUnroll]
31        for (int i = 0; i < N; i++)
32        {
33            values[i] = c;
34        }
35    }
36
37    static Differential dadd(Differential a, Differential b)
38    {
39        Differential output;
40
41        for (int i = 0; i < N; i++)
42        {
43            output.values[i] = T.dadd(a.values[i], b.values[i]);
44        }
45
46        return output;
47    }
48
49    
50    static Differential dmul<U: __BuiltinRealType>(U a, Differential b)
51    {
52        Differential output;
53
54        for (int i = 0; i < N; i++)
55        {
56            output.values[i] = T.dmul<U>(a, b.values[i]);
57        }
58
59        return output;
60    }
61
62    static Differential dzero()
63    {
64        Differential output;
65
66        for (int i = 0; i < N; i++)
67        {
68            output.values[i] = T.dzero();
69        }
70
71        return output;
72    }
73};
74
75[ForwardDifferentiable]
76__generic<T : IDFloat, let N : int>
77myvector<T, N> operator +(myvector<T, N> a, myvector<T, N> b)
78{
79    myvector<T, N> output;
80    [ForceUnroll]
81    for (int i = 0; i < N; i++)
82    {
83        output.values[i] = a.values[i] + b.values[i];
84    }
85    return output;
86}
87
88[ForwardDifferentiable]
89__generic<T : IDFloat, let N : int>
90myvector<T, N> operator *(myvector<T, N> a, myvector<T, N> b)
91{
92    myvector<T, N> output;
93    [ForceUnroll]
94    for (int i = 0; i < N; i++)
95    {
96        output.values[i] = a.values[i] * b.values[i];
97    }
98    return output;
99}
100
101[ForwardDifferentiable]
102__generic<T : IDFloat, let N : int>
103myvector<T, N> operator *(T a, myvector<T, N> b)
104{
105    myvector<T, N> output;
106    [ForceUnroll]
107    for (int i = 0; i < N; i++)
108    {
109        output.values[i] = a * b.values[i];
110    }
111    return output;
112}
113
114__generic<T : IDFloat, let N : int>
115[ForwardDerivative(dot_jvp)]
116T dot(myvector<T, N> a, myvector<T, N> b)
117{
118    T curr = __realCast<T, float>(0.f);
119    [ForceUnroll]
120    for (int i = 0; i < N; i++)
121    {
122        curr = curr + (a.values[i] * b.values[i]);
123    }
124
125    return curr;
126}
127
128__generic<T : IDFloat, let N : int>
129typedef DifferentialPair<myvector<T, N>> dpvector;
130
131__generic<T : IDFloat, let N : int>
132DifferentialPair<T> dot_jvp(dpvector<T, N> a, dpvector<T, N> b)
133{
134    T.Differential curr_d = (T.dzero());
135    T curr_p = __realCast<T, float>(0.f);
136    [ForceUnroll]
137    for (int i = 0; i < N; i++)
138    {
139        curr_p = curr_p + (a.p.values[i] * b.p.values[i]);
140        curr_d = T.dadd(
141                    curr_d, 
142                    T.dadd(
143                        T.dmul<T>(a.p.values[i], b.d.values[i]),
144                        T.dmul<T>(b.p.values[i], a.d.values[i])));
145    }
146
147    return DifferentialPair<T>(curr_p, curr_d);
148}
149
150__generic<let N : int>
151struct lineardvector : IDifferentiable
152{
153    typedef lineardvector<N> Differential;
154
155    myvector<Real, N>.Differential val;
156
157    __init(vector<Real.Differential, N> a)
158    {
159        [ForceUnroll]
160        for (int i = 0; i < N; i++)
161        {
162            val.values[i] = a[i];
163        }
164    }
165
166    // Add a new constructor for dadd() function.
167    __init(Real a[N])
168    {
169        [ForceUnroll]
170        for (int i = 0; i < N; i++)
171        {
172            val.values[i] = a[i];
173        }
174    }
175};
176
177__generic<let N : int>
178struct linearvector : MyLinearArithmeticType, IDifferentiable
179{
180    typedef lineardvector<N> Differential;
181
182    [DerivativeMember(Differential.val)]
183    myvector<Real, N> val;
184
185    [ForwardDifferentiable]
186    static linearvector<N> ladd(linearvector<N> a, linearvector<N> b)
187    {
188        return linearvector<N>(a.val + b.val);
189    }
190
191    [ForwardDifferentiable]
192    static linearvector<N> lmul(linearvector<N> a, linearvector<N> b)
193    {
194        return linearvector<N>(a.val * b.val);
195    }
196
197    [ForwardDifferentiable]
198    static linearvector<N> lscale(float a, linearvector<N> b)
199    {
200        return linearvector<N>(a * b.val);
201    }
202
203    [ForwardDifferentiable]
204    static float ldot(linearvector<N> a, linearvector<N> b)
205    {
206        return dot(a.val, b.val);
207    }
208
209    static Differential dzero()
210    {
211        lineardvector<N> dout;
212        dout.val = myvector<Real, N>.dzero();
213        return dout;
214    }
215
216    static Differential dadd(Differential a, Differential b)
217    {
218        // return { myvector<Real, N>.dadd(a.val, b.val) };
219        //
220        // Above code will not work because
221        // myvector<Real, N>.dadd will return dvector<T.Differential, N> type
222        // while Differential == lineardvector<N> type
223        // and the constructor of lineardvector<N> requires a vector<Real.Differential, N> type
224        // and dvector<T.Differential, N> != vector<Real.Differential, N>, though they have the
225        // same members.
226        //
227        // In our new design, generic will not be C-Style struct anymore.
228        dvector<Real.Differential, N> d = myvector<Real, N>.dadd(a.val, b.val);
229        return {d.values};
230    }
231
232    static Differential dmul<T: __BuiltinRealType>(T a, Differential b)
233    {
234        dvector<Real.Differential, N> d = myvector<Real, N>.dmul<T>(a, b.val);
235        return {d.values};
236    }
237
238    [ForwardDifferentiable]
239    __init(vector<Real, N> a)
240    {
241        [ForceUnroll]
242        for (int i = 0; i < N; i++)
243        {
244            val.values[i] = a[i];
245        }
246    }
247
248    [ForwardDifferentiable]
249    __init(myvector<Real, N> a)
250    {
251        val = a;
252    }
253};
254
255typedef linearvector<3> myfloat3;
256typedef linearvector<4> myfloat4;
257
258typedef lineardvector<3> mydfloat3;
259typedef lineardvector<4> mydfloat4;
260
261typedef DifferentialPair<Real> dpfloat;
262
263[TreatAsDifferentiable]
264interface MyLinearArithmeticType
265{
266    static This ladd(This a, This b);
267    static This lmul(This a, This b);
268    static This lscale(Real a, This b);
269    static Real ldot(This a, This b);
270};
271
272typedef DifferentialPair<myfloat4> dpfloat4;
273typedef DifferentialPair<myfloat3> dpfloat3;
274
275extension float : MyLinearArithmeticType
276{
277    [ForwardDifferentiable]
278    static float ladd(float a, float b)
279    {
280        return a + b;
281    }
282
283    [ForwardDifferentiable]
284    static float lmul(float a, float b)
285    {
286        return a * b;
287    }
288
289    [ForwardDifferentiable]
290    static float lscale(float a, float b)
291    {
292        return a * b;
293    }
294
295    [ForwardDifferentiable]
296    static float ldot(float a, float b)
297    {
298        return a * b;
299    }
300};
301
302typealias MyLinearArithmeticDifferentiableType = IDifferentiable & MyLinearArithmeticType;
303
304__generic<T : MyLinearArithmeticDifferentiableType>
305[ForwardDifferentiable]
306T operator +(T a, T b)
307{
308    return T.ladd(a, b);
309}
310
311__generic<T : MyLinearArithmeticDifferentiableType>
312[ForwardDifferentiable]
313T operator *(T a, T b)
314{
315    return T.lmul(a, b);
316}
317
318__generic<G : MyLinearArithmeticDifferentiableType>
319[ForwardDifferentiable]
320G f(G x)
321{
322    G a = x + x;
323    G b = x * x;
324
325    return a * a + G.lscale((Real)3.0, x); 
326}
327
328
329[numthreads(1, 1, 1)]
330void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
331{
332    {
333        dpfloat dpa = dpfloat(2.0, 1.0);
334        dpfloat4 dpf4 = dpfloat4(myfloat4(float4(1.5, 2.0, 0.5, 1.0)), mydfloat4(float4(0.5, 0.8, 1.6, 2.5)));
335        dpfloat3 dpf3 = dpfloat3(myfloat3(float3(1.0, 3.0, 5.0)), mydfloat3(float3(0.5, 1.5, 2.5)));
336
337        outputBuffer[0] = f(dpa.p);                               // Expect: 22.0
338        outputBuffer[1] = __fwd_diff(f)(dpfloat(2.0, 0.5)).d;          // Expect: 9.5
339        outputBuffer[2] = __fwd_diff(f)(dpf4).d.val.values[3];         // Expect: 27.5
340        outputBuffer[3] = __fwd_diff(f)(dpf3).d.val.values[1];         // Expect: 40.5
341    }
342}