yum-mirror/slang
Making it easier to work with shaders
git clone https://git.yummers.dev/yum-mirror/slang
f02b08490
master
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