MET-35 Getter (#19)

Have the dummy `AllGetter` that grabs all FMBs for any ray.

Author: arijit-dasgupta

genmetaballs/src/cuda/core/fmb.cuh

@@ -7,7 +7,7 @@ struct FMB {
     float3 extent;
 };
 
-template <typename containter_template>
+template <template <typename> class containter_template>
 class FMBs {
 private:
     containter_template<FMB> fmbs_;
@@ -17,4 +17,10 @@ public:
     FMBs(uint32_t size) : fmbs_(size), log_weights_(size) {
         // TODO: set all log_weights_ to 0
     }
+    CUDA_CALLABLE const containter_template<FMB>& get_all_fmbs() const {
+        return fmbs_;
+    }
+    CUDA_CALLABLE const FMB& get_fmb(uint32_t idx) const {
+        return fmbs_[idx];
+    }
 };

genmetaballs/src/cuda/core/forward.cu

@@ -47,7 +47,7 @@ __global__ render_kernel(const Getter fmb_getter, const Blender blender,
 template <class Getter, class Intersector, class Blender, class Confidence>
 void render_fmbs(const FMBs& fmbs, const Intrinsics& intr, const Pose& extr) {
     // initialize the fmb_getter
-    typename Getter::Getter fmb_getter(fmbs, intr, extr);
+    typename Getter::Getter fmb_getter(fmbs, extr);
     auto kernel = render_kernel<Getter, Intersector, Blender, Confidence>;
     kernel<<<NUM_BLOCKS, THREADS_PER_BLOCK>>>(fmb_getter, fmbs, intr, extr);
 }

genmetaballs/src/cuda/core/getter.cuh

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cmath>
+#include <cuda_runtime.h>
+
+#include "camera.cuh"
+#include "fmb.cuh"
+#include "geometry.cuh"
+#include "utils.cuh"
+
+// This is the dummy version of getter, where all FMBs are relevant to any ray
+template <template <typename> class containter_template>
+struct AllGetter {
+    const FMBs<containter_template>& fmbs;
+    Pose extr; // Current assumption: rays are in camera frame
+
+    CUDA_CALLABLE AllGetter(const FMBs<containter_template>& fmbs, const Pose& extr)
+        : fmbs(fmbs), extr(extr) {}
+
+    // It does not bother using the ray, because it simply returns all FMBs
+    CUDA_CALLABLE const containter_template<FMB>& get_metaballs(const Ray& ray) const {
+        return fmbs.get_all_fmbs();
+    }
+};

tests/cpp_tests/test_getter.cu

@@ -0,0 +1,139 @@
+#include <cstdint>
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+#include <type_traits>
+#include <vector>
+
+#include "core/fmb.cuh"
+#include "core/geometry.cuh"
+#include "core/getter.cuh"
+#include "core/utils.cuh"
+
+// a whacky helper to get the container type from the template parameter
+template <typename Container>
+struct template_container_of;
+
+template <template <typename...> class Template, typename... Args>
+struct template_container_of<Template<Args...>> {
+    template <typename T>
+    using type = Template<T>;
+};
+
+template <typename InstantiatedContainer>
+struct GetterTestTypes;
+
+template <template <typename...> class Template, typename... Args>
+struct GetterTestTypes<Template<Args...>> {
+    template <typename T>
+    using ContainerTemplate = typename template_container_of<Template<Args...>>::template type<T>;
+
+    using FMBsType = FMBs<ContainerTemplate>;
+    using AllGetterType = AllGetter<ContainerTemplate>;
+};
+
+// Template fixture for all container types
+template <typename Container>
+class AllGetterTestFixture : public ::testing::Test {};
+
+using ContainerTypes = ::testing::Types<std::vector<FMB>, thrust::device_vector<FMB>>;
+TYPED_TEST_SUITE(AllGetterTestFixture, ContainerTypes);
+
+// CUDA kernel that constructs AllGetter and calls get_metaballs
+template <typename AllGetterType, typename FMBsType>
+__global__ void test_get_metaballs_kernel_device(const FMBsType* fmbs, const Pose* extr,
+                                                 const Ray* rays, int num_rays, int* out_sizes) {
+    int idx = threadIdx.x + blockIdx.x * blockDim.x;
+    AllGetterType getter(*fmbs, *extr);
+    const auto& fmbs_returned = getter.get_metaballs(rays[idx]);
+    out_sizes[idx] = static_cast<int>(fmbs_returned.size());
+}
+
+TYPED_TEST(AllGetterTestFixture, ReturnsAllFMBsForAnyRay) {
+    // Extract types from TypeParam (std::vector<FMB> or thrust::device_vector<FMB>)
+    using FMBsType = typename GetterTestTypes<TypeParam>::FMBsType;
+    using AllGetterType = typename GetterTestTypes<TypeParam>::AllGetterType;
+
+    constexpr uint32_t num_fmbs = 40;
+    FMBsType fmbs(num_fmbs);
+
+    Pose extr = Pose();
+    AllGetterType getter(fmbs, extr);
+
+    // Create test rays
+    std::vector<Ray> rays = {
+        Ray{Vec3D{0.0f, 0.0f, 0.0f}, Vec3D{1.0f, 0.0f, 0.0f}},
+        Ray{Vec3D{1.0f, 1.0f, 1.0f}, Vec3D{0.0f, 1.0f, 0.0f}},
+        Ray{Vec3D{-1.0f, -1.0f, -1.0f}, Vec3D{0.0f, 0.0f, 1.0f}},
+        Ray{Vec3D{2.5f, -3.1f, 0.2f}, Vec3D{-0.5f, 0.6f, 0.0f}},
+        Ray{Vec3D{4.4f, 0.0f, -0.9f}, Vec3D{0.3f, -0.2f, 1.0f}},
+        Ray{Vec3D{5.0f, 2.2f, 1.1f}, Vec3D{-1.0f, 2.0f, 0.2f}},
+        Ray{Vec3D{0.0f, 7.0f, 6.0f}, Vec3D{0.0f, -1.0f, -1.0f}},
+        Ray{Vec3D{-2.0f, 0.0f, 0.0f}, Vec3D{0.2f, 1.1f, 0.7f}},
+        Ray{Vec3D{9.1f, -0.3f, 2.7f}, Vec3D{-0.3f, 0.1f, 0.0f}},
+        Ray{Vec3D{1.2f, 8.8f, -4.5f}, Vec3D{1.0f, 0.0f, 1.0f}},
+    };
+
+    // Get reference to all FMBs from the original FMBs object
+    const auto& all_fmbs_ref = fmbs.get_all_fmbs();
+
+    // Test on host - AllGetter should return the same container for all rays
+    for (const auto& ray : rays) {
+        const auto& fmbs_returned = getter.get_metaballs(ray);
+
+        // Verify that we get the same container reference (both are host objects)
+        EXPECT_EQ(&fmbs_returned, &all_fmbs_ref)
+            << "AllGetter should return the same FMBs container for all rays";
+
+        // For thrust::device_vector, also verify device pointers match
+        if constexpr (std::is_same_v<TypeParam, thrust::device_vector<FMB>>) {
+            EXPECT_EQ(thrust::raw_pointer_cast(fmbs_returned.data()),
+                      thrust::raw_pointer_cast(all_fmbs_ref.data()))
+                << "Device pointers should match for thrust::device_vector";
+        }
+
+        // Verify container sizes match
+        EXPECT_EQ(fmbs_returned.size(), all_fmbs_ref.size())
+            << "Returned FMBs container size must match all_fmbs size";
+    }
+
+    // Test on GPU for device containers
+    // Construct AllGetter on device and call get_metaballs
+    if constexpr (std::is_same_v<TypeParam, thrust::device_vector<FMB>>) {
+        Ray* d_rays = nullptr;
+        FMBsType* d_fmbs = nullptr;
+        Pose* d_extr = nullptr;
+        int* d_sizes = nullptr;
+        int num_rays = static_cast<int>(rays.size());
+
+        CUDA_CHECK(cudaMalloc(&d_rays, num_rays * sizeof(Ray)));
+        CUDA_CHECK(cudaMalloc(&d_fmbs, sizeof(FMBsType)));
+        CUDA_CHECK(cudaMalloc(&d_extr, sizeof(Pose)));
+        CUDA_CHECK(cudaMalloc(&d_sizes, num_rays * sizeof(int)));
+        CUDA_CHECK(cudaMemcpy(d_rays, rays.data(), num_rays * sizeof(Ray), cudaMemcpyHostToDevice));
+        CUDA_CHECK(cudaMemcpy(d_fmbs, &fmbs, sizeof(FMBsType), cudaMemcpyHostToDevice));
+        CUDA_CHECK(cudaMemcpy(d_extr, &extr, sizeof(Pose), cudaMemcpyHostToDevice));
+
+        // Launch kernel that constructs getter and calls get_metaballs on the device
+        test_get_metaballs_kernel_device<AllGetterType, FMBsType>
+            <<<1, num_rays>>>(d_fmbs, d_extr, d_rays, num_rays, d_sizes);
+        CUDA_CHECK(cudaGetLastError());
+        CUDA_CHECK(cudaDeviceSynchronize());
+
+        std::vector<int> host_sizes(num_rays);
+        CUDA_CHECK(
+            cudaMemcpy(host_sizes.data(), d_sizes, num_rays * sizeof(int), cudaMemcpyDeviceToHost));
+
+        // Verify that get_metaballs returns the correct size for all rays
+        for (int i = 0; i < num_rays; ++i) {
+            EXPECT_EQ(static_cast<size_t>(host_sizes[i]), all_fmbs_ref.size())
+                << "Device get_metaballs returned correct size for ray " << i;
+        }
+
+        CUDA_CHECK(cudaFree(d_rays));
+        CUDA_CHECK(cudaFree(d_fmbs));
+        CUDA_CHECK(cudaFree(d_extr));
+        CUDA_CHECK(cudaFree(d_sizes));
+    }
+}