Correct rejection control on GPU

CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.20 FATAL_ERROR)
 project(
   ViennaRay
   LANGUAGES CXX
-  VERSION 3.7.1)
+  VERSION 3.7.2)
 
 # --------------------------------------------------------------------------------------------------------
 # Library switches

README.md

@@ -63,7 +63,7 @@ We recommend using [CPM.cmake](https://github.com/cpm-cmake/CPM.cmake) to consum
 * Installation with CPM
 
   ```cmake
-  CPMAddPackage("gh:viennatools/viennaray@3.7.1") # Use the latest release version
+  CPMAddPackage("gh:viennatools/viennaray@3.7.2") # Use the latest release version
   ```
 
 * With a local installation

gpu/include/raygLaunchParams.hpp

@@ -25,6 +25,7 @@ struct LaunchParams {
   unsigned int numElements;      // to determine result buffer index
   unsigned int *dataPerParticle; // to determine result buffer index
 
+  unsigned maxReflections = std::numeric_limits<unsigned>::max();
   unsigned maxBoundaryHits = 1000;
   uint8_t particleIdx = 0;
   uint8_t particleType = 0;
@@ -68,10 +69,37 @@ getIdxOffset(int dataIdx, const LaunchParams &launchParams) {
 }
 
 __device__ __forceinline__ bool continueRay(const LaunchParams &launchParams,
-                                            const PerRayData &prd) {
-  return prd.rayWeight > launchParams.rayWeightThreshold &&
-         prd.numBoundaryHits < launchParams.maxBoundaryHits &&
-         prd.energy >= 0.f;
+                                            PerRayData &prd) {
+  if (prd.rayWeight <= 0.f || prd.energy < 0.f)
+    return false;
+
+  if (prd.numReflections > launchParams.maxReflections ||
+      prd.numBoundaryHits > launchParams.maxBoundaryHits)
+    return false;
+
+  if (prd.numReflections > 1e4)
+    return false;
+
+  // If the weight of the ray is above a certain threshold, we always reflect.
+  // If the weight of the ray is below the threshold, we randomly decide to
+  // either kill the ray or increase its weight (in an unbiased way).
+  if (prd.rayWeight >= launchParams.rayWeightThreshold && prd.energy >= 0.f)
+    return true;
+
+  // We want to set the weight of (the reflection of) the ray to the value of
+  // renewWeight. In order to stay unbiased we kill the reflection with a
+  // probability of (1 - rayWeight / renewWeight).
+  float renewWeight = 0.3f;
+  float rnd = getNextRand(&prd.RNGstate);
+  float killProbability = 1.f - prd.rayWeight / renewWeight;
+  if (rnd < killProbability) {
+    // kill the ray
+    return false;
+  }
+  // set rayWeight to new weight
+  prd.rayWeight = renewWeight;
+  // continue ray
+  return true;
 }
 #endif
 

gpu/include/raygPerRayData.hpp

@@ -31,6 +31,7 @@ struct PerRayData {
 
   // Hit data
   unsigned int numBoundaryHits = 0;
+  unsigned int numReflections = 0;
   unsigned int primID = 0; // primID of closest hit
   float tMin = 1e20f;      // distance to closest hit
 
@@ -58,16 +59,16 @@ static __forceinline__ __device__ void packPointer(void *ptr, uint32_t &i0,
   i1 = uptr & 0x00000000ffffffff;
 }
 
-template <typename T> static __forceinline__ __device__ T *getPRD() {
+static __forceinline__ __device__ PerRayData *getPRD() {
   const uint32_t u0 = optixGetPayload_0();
   const uint32_t u1 = optixGetPayload_1();
-  return reinterpret_cast<T *>(unpackPointer(u0, u1));
+  return reinterpret_cast<PerRayData *>(unpackPointer(u0, u1));
 }
 
 static __device__ void initializeRNGState(PerRayData *prd,
                                           unsigned int linearLaunchIndex,
                                           unsigned int seed) {
-  auto rngSeed = tea<4>(linearLaunchIndex, seed);
+  auto rngSeed = tea<3>(linearLaunchIndex, seed);
   curand_init(rngSeed, 0, 0, &prd->RNGstate);
 }
 #endif

gpu/include/raygReflection.hpp

@@ -28,7 +28,7 @@ __device__ __inline__ Vec3Df computeNormal(const void *sbtData,
     const Vec3Df &A = sbt->vertex[index[0]];
     const Vec3Df &B = sbt->vertex[index[1]];
     const Vec3Df &C = sbt->vertex[index[2]];
-    return Normalize<float, 3>(CrossProduct<float>(B - A, C - A));
+    return Normalize(CrossProduct(B - A, C - A));
   } break;
   case 1: {
     // Disks

gpu/include/raygTrace.hpp

@@ -107,6 +107,9 @@ template <class T, int D> class Trace {
     // Threshold value for neighbor detection in disk-based geometries
     launchParams.tThreshold = 1.1 * gridDelta_; // TODO: find the best value
 
+    launchParams.maxReflections = config_.maxReflections;
+    launchParams.maxBoundaryHits = config_.maxBoundaryHits;
+
     int numPointsPerDim =
         static_cast<int>(std::sqrt(static_cast<T>(launchParams.numElements)));
 

gpu/pipelines/GeneralPipelineDisk.cu

@@ -22,7 +22,7 @@ extern "C" __constant__ viennaray::gpu::LaunchParams launchParams;
 extern "C" __global__ void __intersection__() {
   const HitSBTDataDisk *sbtData =
       (const HitSBTDataDisk *)optixGetSbtDataPointer();
-  PerRayData *prd = (PerRayData *)getPRD<PerRayData>();
+  PerRayData *prd = getPRD();
 
   // Get the index of the AABB box that was hit
   const unsigned int primID = optixGetPrimitiveIndex();
@@ -69,7 +69,7 @@ extern "C" __global__ void __intersection__() {
 extern "C" __global__ void __closesthit__() {
   const HitSBTDataDisk *sbtData =
       (const HitSBTDataDisk *)optixGetSbtDataPointer();
-  PerRayData *prd = (PerRayData *)getPRD<PerRayData>();
+  PerRayData *prd = getPRD();
 
   const unsigned int primID = optixGetPrimitiveIndex();
   prd->tMin = optixGetRayTmax() - launchParams.tThreshold;
@@ -143,7 +143,7 @@ extern "C" __global__ void __closesthit__() {
   }
 }
 
-extern "C" __global__ void __miss__() { getPRD<PerRayData>()->rayWeight = 0.f; }
+extern "C" __global__ void __miss__() { getPRD()->rayWeight = 0.f; }
 
 extern "C" __global__ void __raygen__() {
   const uint3 idx = optixGetLaunchIndex();
@@ -188,5 +188,6 @@ extern "C" __global__ void __raygen__() {
                0,                             // missSBTIndex
                u0, u1);                       // Payload
     prd.totalCount = 0;                       // Reset PerRayData
+    prd.numReflections++;
   }
 }

gpu/pipelines/GeneralPipelineLine.cu

@@ -22,7 +22,7 @@ extern "C" __constant__ viennaray::gpu::LaunchParams launchParams;
 extern "C" __global__ void __intersection__() {
   const HitSBTDataLine *sbtData =
       (const HitSBTDataLine *)optixGetSbtDataPointer();
-  PerRayData *prd = (PerRayData *)getPRD<PerRayData>();
+  PerRayData *prd = getPRD();
 
   // Get the index of the AABB box that was hit
   const int primID = optixGetPrimitiveIndex();
@@ -50,7 +50,7 @@ extern "C" __global__ void __intersection__() {
 extern "C" __global__ void __closesthit__() {
   const HitSBTDataLine *sbtData =
       (const HitSBTDataLine *)optixGetSbtDataPointer();
-  PerRayData *prd = (PerRayData *)getPRD<PerRayData>();
+  PerRayData *prd = getPRD();
 
   const unsigned int primID = optixGetPrimitiveIndex();
   prd->tMin = optixGetRayTmax();
@@ -80,7 +80,7 @@ extern "C" __global__ void __closesthit__() {
   }
 }
 
-extern "C" __global__ void __miss__() { getPRD<PerRayData>()->rayWeight = 0.f; }
+extern "C" __global__ void __miss__() { getPRD()->rayWeight = 0.f; }
 
 extern "C" __global__ void __raygen__() {
   const uint3 idx = optixGetLaunchIndex();
@@ -125,5 +125,6 @@ extern "C" __global__ void __raygen__() {
                0,                             // missSBTIndex
                u0, u1);                       // Payload
     prd.totalCount = 0;                       // Reset PerRayData
+    prd.numReflections++;
   }
 }

gpu/pipelines/GeneralPipelineTriangle.cu

@@ -24,7 +24,7 @@ extern "C" __constant__ viennaray::gpu::LaunchParams launchParams;
 extern "C" __global__ void __closesthit__() {
   const HitSBTDataTriangle *sbtData =
       (const HitSBTDataTriangle *)optixGetSbtDataPointer();
-  PerRayData *prd = (PerRayData *)getPRD<PerRayData>();
+  PerRayData *prd = getPRD();
 
   const unsigned int primID = optixGetPrimitiveIndex();
   prd->tMin = optixGetRayTmax();
@@ -55,7 +55,7 @@ extern "C" __global__ void __closesthit__() {
   }
 }
 
-extern "C" __global__ void __miss__() { getPRD<PerRayData>()->rayWeight = 0.f; }
+extern "C" __global__ void __miss__() { getPRD()->rayWeight = 0.f; }
 
 extern "C" __global__ void __raygen__() {
   const uint3 idx = optixGetLaunchIndex();
@@ -99,7 +99,7 @@ extern "C" __global__ void __raygen__() {
                1,                             // SBT stride
                0,                             // missSBTIndex
                u0, u1);                       // Payload
-
+    prd.numReflections++;
 #ifdef COUNT_RAYS
     int *counter = reinterpret_cast<int *>(launchParams.customData);
     atomicAdd(counter, 1);

include/viennaray/rayTraceKernel.hpp

@@ -211,7 +211,7 @@ template <typename NumericType, int D, GeometryType geoType> class TraceKernel {
 
           /* -------- Boundary hit -------- */
           if (rayHit.hit.geomID == boundaryID) {
-            if (++boundaryHits > 1000) {
+            if (++boundaryHits > config_.maxBoundaryHits) {
               // terminate ray if too many boundary hits
               ++raysTerminated;
               break;
@@ -321,10 +321,6 @@ template <typename NumericType, int D, GeometryType geoType> class TraceKernel {
           if (rayWeight <= 0) {
             break;
           }
-          reflect = rejectionControl(rayWeight, initialRayWeight, rngState);
-          if (!reflect) {
-            break;
-          }
           if (++numReflections > config_.maxReflections) {
             // terminate ray if too many reflections
             break;
@@ -334,6 +330,10 @@ template <typename NumericType, int D, GeometryType geoType> class TraceKernel {
             ++raysTerminated;
             break;
           }
+          reflect = rejectionControl(rayWeight, initialRayWeight, rngState);
+          if (!reflect) {
+            break;
+          }
 
           // Update ray direction and origin
           fillRayPosition(rayHit.ray, hitPoint);
@@ -460,9 +460,9 @@ template <typename NumericType, int D, GeometryType geoType> class TraceKernel {
     // We want to set the weight of (the reflection of) the ray to the value of
     // renewWeight. In order to stay unbiased we kill the reflection with a
     // probability of (1 - rayWeight / renewWeight).
-    auto rnd = static_cast<double>(rng() / RNG::max());
+    std::uniform_real_distribution<> dist;
     auto killProbability = 1.0 - rayWeight / renewWeight;
-    if (rnd < killProbability) {
+    if (dist(rng) < killProbability) {
       // kill the ray
       return false;
     }

include/viennaray/rayUtil.hpp

@@ -83,6 +83,7 @@ struct KernelConfig {
   size_t numRaysPerPoint = 1000;
   size_t numRaysFixed = 0;
   unsigned maxReflections = std::numeric_limits<unsigned>::max();
+  unsigned maxBoundaryHits = 1000;
   unsigned rngSeed = 0;
 
   bool useRandomSeed = true;