Implement sub-sampling in raytracing code (#277)

* Implement sub-sampling in raytracing code

* Review fixes

include/rgl/api/core.h

@@ -382,6 +382,15 @@ typedef enum : int32_t
 	RGL_FIELD_DYNAMIC_FORMAT = 13842,
 } rgl_field_t;
 
+/**
+ * Kinds of return type for multi-return LiDAR output.
+ */
+typedef enum : int32_t
+{
+	RGL_RETURN_TYPE_FIRST = 0,
+	RGL_RETURN_TYPE_LAST = 1,
+} rgl_return_type_t;
+
 /**
  * Helper enum for axis selection
  */

src/gpu/MultiReturn.hpp

@@ -0,0 +1,28 @@
+// Copyright 2024 Robotec.AI
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <RGLFields.hpp>
+
+#define MULTI_RETURN_BEAM_VERTICES 8
+#define MULTI_RETURN_BEAM_CIRCLES 2
+#define MULTI_RETURN_BEAM_SAMPLES (1 + (MULTI_RETURN_BEAM_VERTICES * MULTI_RETURN_BEAM_CIRCLES))
+
+struct MultiReturnPointers
+{
+	Field<IS_HIT_I32>::type* isHit;
+	Field<XYZ_VEC3_F32>::type* xyz;
+	Field<DISTANCE_F32>::type* distance;
+};

src/gpu/RaytraceRequestContext.hpp

@@ -16,6 +16,7 @@
 
 #include <optix.h>
 #include <RGLFields.hpp>
+#include <gpu/MultiReturn.hpp>
 
 struct RaytraceRequestContext
 {
@@ -38,7 +39,7 @@ struct RaytraceRequestContext
 	const int* ringIds;
 	size_t ringIdsCount;
 
-	const float* rayTimeOffsets;
+	const float* rayTimeOffsetsMs;
 	size_t rayTimeOffsetsCount;
 
 	OptixTraversableHandle scene;
@@ -61,5 +62,9 @@ struct RaytraceRequestContext
 	Field<ELEVATION_F32>::type* elevation;
 	Field<NORMAL_VEC3_F32>::type* normal;
 	Field<INCIDENT_ANGLE_F32>::type* incidentAngle;
+
+	// Multi-Return
+	float beamHalfDivergence;
+	MultiReturnPointers mrSamples;
 };
 static_assert(std::is_trivially_copyable<RaytraceRequestContext>::value);

src/graph/NodesCore.hpp

@@ -30,6 +30,7 @@
 #include <gpu/nodeKernels.hpp>
 #include <CacheManager.hpp>
 #include <GPUFieldDescBuilder.hpp>
+#include <gpu/MultiReturn.hpp>
 
 
 struct FormatPointsNode : IPointsNodeSingleInput
@@ -143,6 +144,38 @@ struct RaytraceNode : IPointsNode
 
 	std::unordered_map<rgl_field_t, IAnyArray::Ptr> fieldData; // All should be DeviceAsyncArray
 
+	struct MultiReturnFields
+	{
+		MultiReturnFields(StreamBoundObjectsManager& arrayMgr)
+		  : isHit(DeviceAsyncArray<Field<IS_HIT_I32>::type>::create(arrayMgr)),
+		    xyz(DeviceAsyncArray<Field<XYZ_VEC3_F32>::type>::create(arrayMgr)),
+		    distance(DeviceAsyncArray<Field<DISTANCE_F32>::type>::create(arrayMgr))
+		{}
+		void resize(size_t size)
+		{
+			isHit->resize(size, false, false);
+			xyz->resize(size, false, false);
+			distance->resize(size, false, false);
+		}
+		MultiReturnPointers getPointers()
+		{
+			return MultiReturnPointers{
+			    .isHit = isHit->getWritePtr(),
+			    .xyz = xyz->getWritePtr(),
+			    .distance = distance->getWritePtr(),
+			};
+		}
+		DeviceAsyncArray<Field<IS_HIT_I32>::type>::Ptr isHit;
+		DeviceAsyncArray<Field<XYZ_VEC3_F32>::type>::Ptr xyz;
+		DeviceAsyncArray<Field<DISTANCE_F32>::type>::Ptr distance;
+	};
+
+
+	MultiReturnFields mrSamples = MultiReturnFields{arrayMgr};
+	MultiReturnFields mrFirst = MultiReturnFields{arrayMgr};
+	MultiReturnFields mrLast = MultiReturnFields{arrayMgr};
+
+
 	template<rgl_field_t>
 	auto getPtrTo();
 

src/graph/RaytraceNode.cpp

@@ -65,6 +65,11 @@ void RaytraceNode::enqueueExecImpl()
 		data->resize(raysNode->getRayCount(), false, false);
 	}
 
+	// TODO(prybicki): don't update this when not needed
+	mrSamples.resize(MULTI_RETURN_BEAM_SAMPLES * raysNode->getRayCount());
+	mrFirst.resize(raysNode->getRayCount());
+	mrLast.resize(raysNode->getRayCount());
+
 	// Even though we are in graph thread here, we can access Scene class (see comment there)
 	const Mat3x4f* raysPtr = raysNode->getRays()->asSubclass<DeviceAsyncArray>()->getReadPtr();
 	auto sceneAS = Scene::instance().getASLocked();
@@ -92,7 +97,7 @@ void RaytraceNode::enqueueExecImpl()
 	    .rayRangesCount = rayRanges.has_value() ? (*rayRanges)->getCount() : defaultRange->getCount(),
 	    .ringIds = ringIds.has_value() ? (*ringIds)->asSubclass<DeviceAsyncArray>()->getReadPtr() : nullptr,
 	    .ringIdsCount = ringIds.has_value() ? (*ringIds)->getCount() : 0,
-	    .rayTimeOffsets = timeOffsets.has_value() ? (*timeOffsets)->asSubclass<DeviceAsyncArray>()->getReadPtr() : nullptr,
+	    .rayTimeOffsetsMs = timeOffsets.has_value() ? (*timeOffsets)->asSubclass<DeviceAsyncArray>()->getReadPtr() : nullptr,
 	    .rayTimeOffsetsCount = timeOffsets.has_value() ? (*timeOffsets)->getCount() : 0,
 	    .scene = sceneAS,
 	    .sceneTime = Scene::instance().getTime().value_or(Time::zero()).asSeconds(),
@@ -112,6 +117,8 @@ void RaytraceNode::enqueueExecImpl()
 	    .elevation = getPtrTo<ELEVATION_F32>(),
 	    .normal = getPtrTo<NORMAL_VEC3_F32>(),
 	    .incidentAngle = getPtrTo<INCIDENT_ANGLE_F32>(),
+	    .beamHalfDivergence = 0.0f, // TODO: provide API to set externally
+	    .mrSamples = mrSamples.getPointers(),
 	};
 
 	requestCtxDev->copyFrom(requestCtxHst);

test/CMakeLists.txt

@@ -48,6 +48,7 @@ set(RGL_TEST_FILES
     src/synchronization/graphThreadSynchronization.cpp
     src/synchronization/testKernel.cu
     src/scene/incidentAngleTest.cpp
+    src/graph/multiReturnTest.cpp
 )
 
 if (RGL_BUILD_ROS2_EXTENSION)

test/src/graph/multiReturnTest.cpp

@@ -0,0 +1,40 @@
+#include <helpers/commonHelpers.hpp>
+#include "math/Mat3x4f.hpp"
+#include "helpers/lidarHelpers.hpp"
+#include "RGLFields.hpp"
+#include "helpers/testPointCloud.hpp"
+#include "helpers/sceneHelpers.hpp"
+
+class GraphMultiReturn : public RGLTest
+{};
+
+TEST_F(GraphMultiReturn, basic)
+{
+	rgl_node_t useRays = nullptr, raytrace = nullptr, lidarPose = nullptr, compact = nullptr, yield = nullptr;
+	std::vector<rgl_field_t> mrFields = {IS_HIT_I32, DISTANCE_F32, XYZ_VEC3_F32, INTENSITY_F32, ENTITY_ID_I32};
+
+	std::vector<rgl_mat3x4f> rays = {Mat3x4f::identity().toRGL()};
+	rgl_mat3x4f lidarPoseTf = Mat3x4f::TRS({1, 2, 3}, {0, 30, 0}).toRGL();
+
+	rgl_entity_t cube = makeEntity();
+	rgl_mat3x4f cubeTf = Mat3x4f::TRS({0, 0, 0}, {0, 0, 0}, {10, 10, 10}).toRGL();
+	EXPECT_RGL_SUCCESS(rgl_entity_set_pose(cube, &cubeTf));
+
+	EXPECT_RGL_SUCCESS(rgl_node_rays_from_mat3x4f(&useRays, rays.data(), rays.size()));
+	EXPECT_RGL_SUCCESS(rgl_node_rays_transform(&lidarPose, &lidarPoseTf));
+	EXPECT_RGL_SUCCESS(rgl_node_raytrace(&raytrace, nullptr));
+	EXPECT_RGL_SUCCESS(rgl_node_points_compact_by_field(&compact, RGL_FIELD_IS_HIT_I32));
+	EXPECT_RGL_SUCCESS(rgl_node_points_yield(&yield, mrFields.data(), mrFields.size()));
+
+	EXPECT_RGL_SUCCESS(rgl_graph_node_add_child(useRays, lidarPose));
+	EXPECT_RGL_SUCCESS(rgl_graph_node_add_child(lidarPose, raytrace));
+	EXPECT_RGL_SUCCESS(rgl_graph_node_add_child(raytrace, compact));
+	EXPECT_RGL_SUCCESS(rgl_graph_node_add_child(compact, yield));
+
+	EXPECT_RGL_SUCCESS(rgl_graph_run(useRays));
+
+	TestPointCloud pc = TestPointCloud::createFromNode(yield, mrFields);
+	EXPECT_EQ(pc.getPointCount(), 1);
+
+	EXPECT_RGL_SUCCESS(rgl_graph_destroy(useRays));
+}