yum-mirror/slang

Making it easier to work with shaders

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

James Helferty (NVIDIA)render-test: Change D3D12 default to sm_6_5 (#8320)f02b08490

master
8.2 KiB299 linesraw
1//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -shaderobj -output-using-type -profile cs_5_1 -dx12 -use-dxbc -compute-dispatch 4,1,1
2//TEST(compute):COMPARE_COMPUTE_EX(filecheck-buffer=CHECK):-slang -compute -shaderobj -output-using-type -profile cs_5_1 -cuda -use-dxbc -compute-dispatch 4,1,1
3
4//TEST_INPUT:ubuffer(data=[0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0], stride=4):out,name=endpointDifferentialBuffer
5RWStructuredBuffer<float> endpointDifferentialBuffer;
6
7//TEST_INPUT:ubuffer(data=[0 0 0 0 0 0 0 0 0 0], stride=4):out,name=endpointDifferentialBufferInt
8RWStructuredBuffer<int> endpointDifferentialBufferInt;
9
10//TEST_INPUT:ubuffer(data=[0.3 0.7 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0], stride=4):name=endpointBuffer
11RWStructuredBuffer<float> endpointBuffer;
12//TEST_INPUT:ubuffer(data=[1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0], stride=4):name=colorBuffer
13RWStructuredBuffer<float> colorBuffer;
14
15typedef float Color;
16
17struct PRNG
18{
19    __init(uint seed)
20    {
21        this.state = seed;
22    }
23
24    [mutating] uint next()
25    {
26        state ^= state << 13;
27        state ^= state >> 7;
28        state ^= state << 17;
29        return state;
30    }
31
32    [mutating] float nextFloat1D()
33    {
34        return float(next()) / float(4294967295.0);
35    }
36
37    uint state;
38};
39
40struct LineSegment : IDifferentiable
41{
42    float x0;
43    float x1;
44
45    Color color;
46
47    [BackwardDifferentiable]
48    __init(float _x0, float _x1, Color _color)
49    {
50        x0 = _x0;
51        x1 = _x1;
52        color = _color;
53    }
54};
55
56struct Intersection : IDifferentiable
57{
58    LineSegment ls;
59    float x;
60    bool isIntersected;
61    float wt;
62
63    [BackwardDifferentiable]
64    __init(LineSegment _ls, float _x, bool _isIntersected, float _wt)
65    {
66        this.ls = _ls;
67        this.x = _x;
68        this.isIntersected = _isIntersected;
69        this.wt = _wt;
70    }
71};
72
73[BackwardDerivative(d_loadLineSegment)]
74[ForwardDerivative(fwd_loadLineSegment)]
75LineSegment loadLineSegment(uint id)
76{
77    return {endpointBuffer[id * 2], endpointBuffer[id * 2 + 1], colorBuffer[id]};
78}
79
80[BackwardDerivative(d_fwd_loadLineSegment)]
81DifferentialPair<LineSegment> fwd_loadLineSegment(uint id)
82{
83    return DifferentialPair<LineSegment>(loadLineSegment(id), LineSegment.dzero());
84}
85
86void accumulateDifferentialFixedPoint(
87    RWStructuredBuffer<int> buffer, 
88    uint index,
89    float.Differential df,
90    float scale = 1000000.f)
91{
92    InterlockedAdd(buffer[index], (int)round(df * scale));
93}
94
95void d_loadLineSegment(uint id, LineSegment.Differential d_ls)
96{
97    accumulateDifferentialFixedPoint(endpointDifferentialBufferInt, id * 2, d_ls.x0);
98    accumulateDifferentialFixedPoint(endpointDifferentialBufferInt, id * 2 + 1, d_ls.x1);
99}
100
101void d_fwd_loadLineSegment(uint id, DifferentialPair<LineSegment>.Differential dp_ls)
102{
103    accumulateDifferentialFixedPoint(endpointDifferentialBufferInt, id * 2, dp_ls.p.x0);
104    accumulateDifferentialFixedPoint(endpointDifferentialBufferInt, id * 2 + 1, dp_ls.p.x1);
105}
106
107int getIntersectionID(float x)
108{
109    // Line segments are ordered by z-index so return the first intersection.
110    for (int id = 0; id < 2; id++)
111    {
112        LineSegment ls = loadLineSegment(id);
113        if (x > ls.x0 && x < ls.x1)
114            return id;
115    }
116    return -1;
117}
118
119[BackwardDifferentiable]
120Intersection intersect(float x)
121{
122    int id = getIntersectionID(x);
123    if (id >= 0)
124        return Intersection(loadLineSegment((uint)id), x, true, 1.0);
125
126    return Intersection(LineSegment(0, 0, 0), x, false, 0.0);
127}
128
129[BackwardDifferentiable]
130float shadeIntersection(Intersection isect)
131{
132    return isect.ls.color;
133}
134
135float sample1DNormal(inout PRNG prng, float mu, float sigma)
136{
137    float u = prng.nextFloat1D();
138    float v = prng.nextFloat1D();
139    return mu + (sqrt(-2 * log(u))*cos(2*3.1415*v) * sigma);
140}
141
142[BackwardDifferentiable]
143float pdf1DNormal(no_diff float x, float mu, no_diff float sigma)
144{  
145    float k = ((x - mu) / sigma);
146    return exp(-0.5 * (k * k)) / (sigma * 2.506628);
147}
148
149float boundaryTerm(Intersection isect)
150{
151    if (!isect.isIntersected)
152        return 100.0; // Large default value for missed rays.
153    
154    float leftDist = abs(isect.x - isect.ls.x0);
155    float rightDist = abs(isect.ls.x1 - isect.x);
156
157    if (leftDist > rightDist)
158        return rightDist * 30.f;
159    else
160        return leftDist * 30.f;
161}
162
163[BackwardDifferentiable]
164DifferentialPair<float> infinitesimal(DifferentialPair<float> x)
165{
166    return diffPair(x.p - detach(x.p), x.d - detach(x.d));
167}
168
169[BackwardDifferentiable]
170float harmonicWeight(Intersection isect, no_diff Intersection aux_isect)
171{
172    float x_dist = isect.x - aux_isect.x;
173    float k = 1.0 / (((x_dist * x_dist) + no_diff(boundaryTerm(aux_isect))));
174    return k;
175}
176
177[BackwardDifferentiable]
178float attachToGeometry(Intersection isect)
179{
180    float leftWt = detach(isect.ls.x1 - isect.x);
181    float rightWt = detach(isect.x - isect.ls.x0);
182
183    return (leftWt * isect.ls.x0 + rightWt * isect.ls.x1) / (leftWt + rightWt);
184}
185
186[BackwardDifferentiable]
187float warp(Intersection isect, inout PRNG prng)
188{
189    float totalWeight = 0.f;
190    float totalWarpedPoint = 0.f;
191
192    float aux_sigma = 0.01;
193
194    for (int i = 0; i < 32; i++)
195    {
196        float y = no_diff(sample1DNormal(prng, isect.x, aux_sigma));
197        float y_flipped = 2 * isect.x - y;
198
199        Intersection aux_isect_left = intersect(y);
200
201        if (aux_isect_left.isIntersected)
202        {
203            float pdf = pdf1DNormal(y, isect.x, aux_sigma);
204            float wt = harmonicWeight(isect, aux_isect_left) * (pdf / detach(pdf));
205            totalWarpedPoint += attachToGeometry(aux_isect_left) * wt;
206            totalWeight += wt;
207        }
208        
209        Intersection aux_isect_right = intersect(detach(y_flipped));
210
211        if (aux_isect_right.isIntersected)
212        {
213            float pdf = pdf1DNormal(y_flipped, isect.x, aux_sigma);
214            float wt = harmonicWeight(isect, aux_isect_right) * (pdf / detach(pdf));
215            totalWarpedPoint += attachToGeometry(aux_isect_right) * wt;
216            totalWeight += wt;
217        }
218    }
219
220    return totalWarpedPoint / totalWeight;
221}
222
223[BackwardDifferentiable]
224Intersection warpedIntersect(float x, inout PRNG prng)
225{   
226    // TODO: For now the jacobian here is 1.0, 
227    // but we will need to adjust the warp by the jacobian for
228    // more complex intersection models.
229    //
230    Intersection isect = intersect(x);
231
232    Intersection.Differential d_isect = Intersection.Differential.dzero();
233    d_isect.x = 1.0;
234
235    var dpwarp = infinitesimal(
236        __fwd_diff(warp)(diffPair(isect, d_isect), prng));
237
238    isect.x = detach(isect.x) + dpwarp.p;
239    isect.wt = isect.wt * (1 + dpwarp.d);
240
241    return isect;
242}
243
244[BackwardDifferentiable]
245float renderSample(inout PRNG prng) 
246{
247    float u = no_diff(prng.nextFloat1D());
248    
249    float leftBound = 0.0;
250    float rightBound = 1.0;
251
252    float sample = leftBound * u + rightBound * (1 - u);
253    float weight = 1.0/(rightBound - leftBound);
254 
255    Intersection isect = warpedIntersect(sample, prng);
256
257    return shadeIntersection(isect) * isect.wt;
258}
259
260[numthreads(256, 1, 1)]
261void computeMain(uint3 threadIdx : SV_DispatchThreadID,)
262{
263    uint seed = (threadIdx.x * threadIdx.x) * 30 + 3;
264    PRNG prng = PRNG(seed);
265    
266    float d_color = 1.0 / 1000.0;
267    __bwd_diff(renderSample)(prng, d_color);
268
269    AllMemoryBarrierWithGroupSync();
270    
271    // Convert to floating point (but with 2 fewer digits of precision to 
272    // avoid platform-specific differences in floating point precision)
273    // 
274    if (threadIdx.x < 10)
275        endpointDifferentialBuffer[threadIdx.x] =
276            ((endpointDifferentialBufferInt[threadIdx.x]/1000) / 1000000.f) * 1000.f;
277
278// Note that this specific derivative estimation method is biased, so the
279// expected results are approximate. (We've fixed the RNG seed to generate
280// repeatable results)
281//
282// Expect: Approximately -1.0 in endpointDifferentialBuffer[0]
283// Expect: Approximately 1.0 in endpointDifferentialBuffer[1]
284//
285// Expect: Approximately 0.0 in endpointDifferentialBuffer[2]
286// Expect: Approximately 0.0 in endpointDifferentialBuffer[3]
287//
288}
289// CHECK: type: float
290// CHECK-NEXT: -0.{{9[5-9][0-9]}}000
291// CHECK-NEXT: 0.{{9[5-9][0-9]}}000
292// CHECK-NEXT: 0.000000
293// CHECK-NEXT: 0.004000
294// CHECK-NEXT: 0.000000
295// CHECK-NEXT: 0.000000
296// CHECK-NEXT: 0.000000
297// CHECK-NEXT: 0.000000
298// CHECK-NEXT: 0.000000
299// CHECK-NEXT: 0.000000