Architectural changes to support multi-return feature (#285)

* Architectural changes to support multi-return feature

src/api/apiCore.cpp

@@ -1311,6 +1311,27 @@ void TapeCore::tape_node_gaussian_noise_distance(const YAML::Node& yamlNode, Pla
 	state.nodes.insert({nodeId, node});
 }
 
+RGL_API rgl_status_t rgl_node_multi_return_switch(rgl_node_t* node, rgl_return_type_t return_type)
+{
+	auto status = rglSafeCall([&]() {
+		RGL_API_LOG("rgl_node_multi_return_switch(node={}, return_type={})", repr(node), return_type);
+		CHECK_ARG(node != nullptr);
+
+		createOrUpdateNode<MultiReturnSwitchNode>(node, return_type);
+	});
+	TAPE_HOOK(node, return_type);
+	return status;
+}
+
+void TapeCore::tape_node_multi_return_switch(const YAML::Node& yamlNode, PlaybackState& state)
+{
+	auto nodeId = yamlNode[0].as<TapeAPIObjectID>();
+	auto return_type = static_cast<rgl_return_type_t>(yamlNode[1].as<int>());
+	rgl_node_t node = state.nodes.contains(nodeId) ? state.nodes.at(nodeId) : nullptr;
+	rgl_node_multi_return_switch(&node, return_type);
+	state.nodes.insert({nodeId, node});
+}
+
 rgl_status_t rgl_node_is_alive(rgl_node_t node, bool* out_alive)
 {
 	auto status = rglSafeCall([&]() {

src/gpu/nodeKernels.cu

@@ -229,25 +229,30 @@ __global__ void kProcessBeamSamplesFirstLast(size_t beamCount, int samplesPerBea
 	LIMIT(beamCount);
 
 	const auto beamIdx = tid;
-	int firstIdx = 0;
-	int lastIdx = 0;
+	int firstIdx = -1;
+	int lastIdx = -1;
 	for (int sampleIdx = 0; sampleIdx < samplesPerBeam; ++sampleIdx) {
 		if (beamSamples.isHit[beamIdx * samplesPerBeam + sampleIdx] == 0) {
 			continue;
 		}
-		if (beamSamples.distance[beamIdx * samplesPerBeam + sampleIdx] <
-		    beamSamples.distance[beamIdx * samplesPerBeam + firstIdx]) {
+		auto currentFirstDistance = firstIdx >= 0 ? beamSamples.distance[beamIdx * samplesPerBeam + firstIdx] : FLT_MAX;
+		auto currentLastDistance = lastIdx >= 0 ? beamSamples.distance[beamIdx * samplesPerBeam + lastIdx] : -FLT_MAX;
+		if (beamSamples.distance[beamIdx * samplesPerBeam + sampleIdx] < currentFirstDistance) {
 			firstIdx = sampleIdx;
 		}
-		if (beamSamples.distance[beamIdx * samplesPerBeam + sampleIdx] >
-		    beamSamples.distance[beamIdx * samplesPerBeam + lastIdx]) {
+		if (beamSamples.distance[beamIdx * samplesPerBeam + sampleIdx] > currentLastDistance) {
 			lastIdx = sampleIdx;
 		}
 	}
-	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];
+	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];
+	}
 }
 
 
@@ -327,4 +332,4 @@ void gpuProcessBeamSamplesFirstLast(cudaStream_t stream, size_t beamCount, int s
                                     MultiReturnPointers first, MultiReturnPointers last)
 {
 	run(kProcessBeamSamplesFirstLast, stream, beamCount, samplesPerBeam, beamSamples, first, last);
-}
+}

src/graph/NodesCore.hpp

@@ -128,6 +128,8 @@ struct RaytraceNode : IPointsNode
 		return std::const_pointer_cast<const IAnyArray>(fieldData.at(field));
 	}
 
+	IAnyArray::ConstPtr getFieldDataMultiReturn(rgl_field_t field, rgl_return_type_t);
+
 	// RaytraceNode specific
 	void setVelocity(const Vec3f& linearVelocity, const Vec3f& angularVelocity);
 	void enableRayDistortion(bool enabled) { doApplyDistortion = enabled; }
@@ -193,6 +195,31 @@ struct RaytraceNode : IPointsNode
 	void setFields(const std::set<rgl_field_t>& fields);
 };
 
+struct MultiReturnSwitchNode : IPointsNodeSingleInput
+{
+	using Ptr = std::shared_ptr<MultiReturnSwitchNode>;
+	void setParameters(rgl_return_type_t returnType) { this->returnType = returnType; }
+
+	// Node
+	void validateImpl() override
+	{
+		IPointsNodeSingleInput::validateImpl();
+		rtxInput = getExactlyOneInputOfType<RaytraceNode>();
+	}
+
+	void enqueueExecImpl() override {}
+
+	// Data getters
+	IAnyArray::ConstPtr getFieldData(rgl_field_t field) override
+	{
+		return rtxInput->getFieldDataMultiReturn(field, returnType);
+	}
+
+private:
+	rgl_return_type_t returnType;
+	RaytraceNode::Ptr rtxInput;
+};
+
 struct TransformPointsNode : IPointsNodeSingleInput
 {
 	using Ptr = std::shared_ptr<TransformPointsNode>;

src/graph/RaytraceNode.cpp

@@ -137,6 +137,27 @@ void RaytraceNode::enqueueExecImpl()
 	                               mrSamples.getPointers(), mrFirst.getPointers(), mrLast.getPointers());
 }
 
+IAnyArray::ConstPtr RaytraceNode::getFieldDataMultiReturn(rgl_field_t field, rgl_return_type_t type)
+{
+	if (type == RGL_RETURN_TYPE_FIRST) {
+		switch (field) {
+			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)));
+		}
+	}
+	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)));
+		}
+	}
+	throw InvalidPipeline(fmt::format("Unknown multi-return type ({})", type));
+}
+
 void RaytraceNode::setFields(const std::set<rgl_field_t>& fields)
 {
 	auto keyViewer = std::views::keys(fieldData);

src/tape/TapeCore.hpp

@@ -68,6 +68,7 @@ class TapeCore
 	static void tape_node_gaussian_noise_angular_ray(const YAML::Node& yamlNode, PlaybackState& state);
 	static void tape_node_gaussian_noise_angular_hitpoint(const YAML::Node& yamlNode, PlaybackState& state);
 	static void tape_node_gaussian_noise_distance(const YAML::Node& yamlNode, PlaybackState& state);
+	static void tape_node_multi_return_switch(const YAML::Node& yamlNode, PlaybackState& state);
 
 	// Called once in the translation unit
 	static inline bool autoExtendTapeFunctions = std::invoke([]() {
@@ -123,6 +124,7 @@ class TapeCore
 		    TAPE_CALL_MAPPING("rgl_node_gaussian_noise_angular_ray", TapeCore::tape_node_gaussian_noise_angular_ray),
 		    TAPE_CALL_MAPPING("rgl_node_gaussian_noise_angular_hitpoint", TapeCore::tape_node_gaussian_noise_angular_hitpoint),
 		    TAPE_CALL_MAPPING("rgl_node_gaussian_noise_distance", TapeCore::tape_node_gaussian_noise_distance),
+		    TAPE_CALL_MAPPING("rgl_node_multi_return_switch", TapeCore::tape_node_multi_return_switch),
 		};
 		TapePlayer::extendTapeFunctions(tapeFunctions);
 		return true;