Compute XYZ in MR based on distance, not XYZ samples (#295)

* Add RGL_RETURN_TYPE_NOT_DIVERGENT

* Compute XYZ based on raydir and distance

* Fix XYZ computation, run beam sample processing conditionally

include/rgl/api/core.h

@@ -394,8 +394,9 @@ typedef enum : int32_t
  */
 typedef enum : int32_t
 {
-	RGL_RETURN_TYPE_FIRST = 0,
-	RGL_RETURN_TYPE_LAST = 1,
+	RGL_RETURN_TYPE_NOT_DIVERGENT = 0,
+	RGL_RETURN_TYPE_FIRST = 1,
+	RGL_RETURN_TYPE_LAST = 2,
 } rgl_return_type_t;
 
 /**

src/gpu/RaytraceRequestContext.hpp

@@ -28,7 +28,7 @@ struct RaytraceRequestContext
 	float nearNonHitDistance;
 	float farNonHitDistance;
 
-	const Mat3x4f* rays;
+	const Mat3x4f* raysWorld;
 	size_t rayCount;
 
 	Mat3x4f rayOriginToWorld;

src/gpu/nodeKernels.cu

@@ -224,7 +224,7 @@ __global__ void kFilterGroundPoints(size_t pointCount, const Vec3f sensor_up_vec
 }
 
 __global__ void kProcessBeamSamplesFirstLast(size_t beamCount, int samplesPerBeam, MultiReturnPointers beamSamples,
-                                             MultiReturnPointers first, MultiReturnPointers last)
+                                             MultiReturnPointers first, MultiReturnPointers last, const Mat3x4f* beamsWorld)
 {
 	LIMIT(beamCount);
 
@@ -244,14 +244,16 @@ __global__ void kProcessBeamSamplesFirstLast(size_t beamCount, int samplesPerBea
 			lastIdx = sampleIdx;
 		}
 	}
+	Vec3f beamOrigin = beamsWorld[beamIdx] * Vec3f{0, 0, 0};
+	Vec3f beamDir = ((beamsWorld[beamIdx] * Vec3f{0, 0, 1}) - beamOrigin).normalized();
 	bool isHit = firstIdx >= 0; // Note that firstHit >= 0 implies lastHit >= 0
 	first.isHit[beamIdx] = isHit;
 	last.isHit[beamIdx] = isHit;
 	if (isHit) {
-		first.xyz[beamIdx] = beamSamples.xyz[beamIdx * samplesPerBeam + firstIdx];
 		first.distance[beamIdx] = beamSamples.distance[beamIdx * samplesPerBeam + firstIdx];
-		last.xyz[beamIdx] = beamSamples.xyz[beamIdx * samplesPerBeam + lastIdx];
 		last.distance[beamIdx] = beamSamples.distance[beamIdx * samplesPerBeam + lastIdx];
+		first.xyz[beamIdx] = beamOrigin + beamDir * first.distance[beamIdx];
+		last.xyz[beamIdx] = beamOrigin + beamDir * last.distance[beamIdx];
 	}
 }
 
@@ -329,7 +331,7 @@ void gpuRadarComputeEnergy(cudaStream_t stream, size_t count, float rayAzimuthSt
 }
 
 void gpuProcessBeamSamplesFirstLast(cudaStream_t stream, size_t beamCount, int samplesPerBeam, MultiReturnPointers beamSamples,
-                                    MultiReturnPointers first, MultiReturnPointers last)
+                                    MultiReturnPointers first, MultiReturnPointers last, const Mat3x4f* beamsWorld)
 {
-	run(kProcessBeamSamplesFirstLast, stream, beamCount, samplesPerBeam, beamSamples, first, last);
+	run(kProcessBeamSamplesFirstLast, stream, beamCount, samplesPerBeam, beamSamples, first, last, beamsWorld);
 }

src/gpu/nodeKernels.hpp

@@ -53,4 +53,4 @@ void gpuRadarComputeEnergy(cudaStream_t stream, size_t count, float rayAzimuthSt
                            const Field<DISTANCE_F32>::type* hitDist, const Field<NORMAL_VEC3_F32>::type* hitNorm,
                            const Field<XYZ_VEC3_F32>::type* hitPos, Vector<3, thrust::complex<float>>* outBUBRFactor);
 void gpuProcessBeamSamplesFirstLast(cudaStream_t stream, size_t beamCount, int samplesPerBeam, MultiReturnPointers beamSamples,
-                                    MultiReturnPointers first, MultiReturnPointers last);
+                                    MultiReturnPointers first, MultiReturnPointers last, const Mat3x4f* beamWorld);

src/gpu/optixPrograms.cu

@@ -56,7 +56,7 @@ extern "C" __global__ void __raygen__()
 		return;
 	}
 
-	Mat3x4f ray = ctx.rays[rayIdx];
+	Mat3x4f ray = ctx.raysWorld[rayIdx];
 	const Mat3x4f rayLocal =
 	    ctx.rayOriginToWorld.inverse() *
 	    ray; // TODO(prybicki): instead of computing inverse, we should pass rays in local CF and then transform them to world CF.
@@ -144,7 +144,9 @@ extern "C" __global__ void __closesthit__()
 		if (!ctx.doApplyDistortion) {
 			return hitWorldRaytraced;
 		}
-		Mat3x4f undistortedRay = ctx.rays[beamIdx] * makeBeamSampleRayTransform(ctx.beamHalfDivergence, circleIdx, vertexIdx);
+		Mat3x4f sampleRayTf = beamSampleRayIdx == 0 ? Mat3x4f::identity() :
+		                                              makeBeamSampleRayTransform(ctx.beamHalfDivergence, circleIdx, vertexIdx);
+		Mat3x4f undistortedRay = ctx.raysWorld[beamIdx] * sampleRayTf;
 		Vec3f undistortedOrigin = undistortedRay * Vec3f{0, 0, 0};
 		Vec3f undistortedDir = undistortedRay * Vec3f{0, 0, 1} - undistortedOrigin;
 		return undistortedOrigin + undistortedDir * distance;
@@ -188,7 +190,6 @@ extern "C" __global__ void __closesthit__()
 
 	// Save sub-sampling results
 	ctx.mrSamples.isHit[mrSamplesIdx] = true;
-	ctx.mrSamples.xyz[mrSamplesIdx] = hitWorldSeenBySensor;
 	ctx.mrSamples.distance[mrSamplesIdx] = distance;
 	if (beamSampleRayIdx != 0) {
 		return;
@@ -269,7 +270,7 @@ __device__ Mat3x4f makeBeamSampleRayTransform(float halfDivergenceAngleRad, unsi
 
 __device__ void saveNonHitRayResult(float nonHitDistance)
 {
-	Mat3x4f ray = ctx.rays[optixGetLaunchIndex().x];
+	Mat3x4f ray = ctx.raysWorld[optixGetLaunchIndex().x];
 	Vec3f origin = ray * Vec3f{0, 0, 0};
 	Vec3f dir = ray * Vec3f{0, 0, 1} - origin;
 	Vec3f displacement = dir.normalized() * nonHitDistance;

src/graph/NodesCore.hpp

@@ -212,6 +212,9 @@ struct MultiReturnSwitchNode : IPointsNodeSingleInput
 	// Data getters
 	IAnyArray::ConstPtr getFieldData(rgl_field_t field) override
 	{
+		if (returnType == RGL_RETURN_TYPE_NOT_DIVERGENT) {
+			return rtxInput->getFieldData(field);
+		}
 		return rtxInput->getFieldDataMultiReturn(field, returnType);
 	}
 

src/graph/RaytraceNode.cpp

@@ -94,7 +94,7 @@ void RaytraceNode::enqueueExecImpl()
 	    .doApplyDistortion = doApplyDistortion,
 	    .nearNonHitDistance = nearNonHitDistance,
 	    .farNonHitDistance = farNonHitDistance,
-	    .rays = raysPtr,
+	    .raysWorld = raysPtr,
 	    .rayCount = raysNode->getRayCount(),
 	    .rayOriginToWorld = raysNode->getCumulativeRayTransfrom(),
 	    .rayRanges = rayRanges.has_value() ? (*rayRanges)->asSubclass<DeviceAsyncArray>()->getReadPtr() :
@@ -134,8 +134,10 @@ void RaytraceNode::enqueueExecImpl()
 	CHECK_OPTIX(optixLaunch(Optix::getOrCreate().pipeline, getStreamHandle(), pipelineArgsPtr, pipelineArgsSize, &sceneSBT,
 	                        launchDims.x, launchDims.y, launchDims.y));
 
-	gpuProcessBeamSamplesFirstLast(getStreamHandle(), raysNode->getRayCount(), MULTI_RETURN_BEAM_SAMPLES,
-	                               mrSamples.getPointers(), mrFirst.getPointers(), mrLast.getPointers());
+	if (beamHalfDivergence > 0.0f) {
+		gpuProcessBeamSamplesFirstLast(getStreamHandle(), raysNode->getRayCount(), MULTI_RETURN_BEAM_SAMPLES,
+		                               mrSamples.getPointers(), mrFirst.getPointers(), mrLast.getPointers(), raysPtr);
+	}
 }
 
 IAnyArray::ConstPtr RaytraceNode::getFieldDataMultiReturn(rgl_field_t field, rgl_return_type_t type)
@@ -145,15 +147,15 @@ IAnyArray::ConstPtr RaytraceNode::getFieldDataMultiReturn(rgl_field_t field, rgl
 			case XYZ_VEC3_F32: return mrFirst.xyz;
 			case DISTANCE_F32: return mrFirst.distance;
 			case IS_HIT_I32: return mrFirst.isHit;
-			default: throw InvalidPipeline(fmt::format("Multi-Return not supported for this field ({})", toString(field)));
+			default: return getFieldData(field);
 		}
 	}
 	if (type == RGL_RETURN_TYPE_LAST) {
 		switch (field) {
 			case XYZ_VEC3_F32: return mrLast.xyz;
 			case DISTANCE_F32: return mrLast.distance;
 			case IS_HIT_I32: return mrLast.isHit;
-			default: throw InvalidPipeline(fmt::format("Multi-Return not supported for this field ({})", toString(field)));
+			default: return getFieldData(field);
 		}
 	}
 	throw InvalidPipeline(fmt::format("Unknown multi-return type ({})", type));