MET-33: Intersector by mugamma · Pull Request #25 · probcomp/GenMetaBalls

mugamma · 2025-11-26T17:54:32Z

No description provided.

linear · 2025-11-26T17:54:35Z

MET-33 Implement `Intersector`

Definition of done:

tests against linear blender in the JAX code pass

Look into mocking for pytest (if i need the data strucs to do something

mugamma

I've added some comments to help with your guys's review. I recommend first merging MET-32: FMB #24 and then moving on to this one, as it is based off of that branch.

mugamma · 2025-11-27T02:57:01Z

genmetaballs/src/cuda/bindings.cu

+    /*
+     * Confidence module bindings
+     */
+
+    nb::module_ confidence = m.def_submodule("confidence");
+    nb::class_<ZeroParameterConfidence>(confidence, "ZeroParameterConfidence")
+        .def(nb::init<>())
+        .def("get_confidence", &ZeroParameterConfidence::get_confidence);
+
+    nb::class_<TwoParameterConfidence>(confidence, "TwoParameterConfidence")
+        .def(nb::init<float, float>())
+        .def("get_confidence", &TwoParameterConfidence::get_confidence);
+


Moved the confidence bindings up so the modules are defined in alphabetical order.

mugamma · 2025-11-27T02:58:08Z

genmetaballs/src/cuda/bindings.cu

+        .def("cov_inv_apply", &FMB::cov_inv_apply,
+             "apply the inverse covariance matrix to the given vector", nb::arg("vec"))


Added a helper method for applying the inverse covariance matrix. This simplifies the implementation of the intersector and the computation of the quadratic form.

mugamma · 2025-11-27T02:59:25Z

genmetaballs/src/cuda/bindings.cu

+            auto [t, d] = LinearIntersector::intersect(fmb, ray, cam_pose);
+            return std::make_tuple(t, d);


This is necessary since nanobind can't bind cuda::std::tuple to a python tuple.

mugamma · 2025-11-27T03:00:00Z

genmetaballs/src/cuda/core/fmb.cu

+CUDA_CALLABLE __forceinline__ Vec3D vecdiv(const Vec3D u, const Vec3D v) {
+    return {u.x / v.x, u.y / v.y, u.z / v.z};
+}


helper function for this file specifically.

maybe I should make this a lambda inside cov_inv_apply?

I'm fine with leaving this here since we're not in the header file anyways.

mugamma · 2025-11-27T03:01:27Z

genmetaballs/src/cuda/core/fmb.cu

+template <>
+__host__ FMBScene<MemoryLocation::HOST>::FMBScene(size_t size)
+    : fmbs_{new FMB[size]}, log_weights_{new float[size]}, size_{size} {}
+
+template <>
+__host__ FMBScene<MemoryLocation::DEVICE>::FMBScene(size_t size) : size_{size} {
+    CUDA_CHECK(cudaMalloc(&fmbs_, size * sizeof(FMB)));
+    CUDA_CHECK(cudaMalloc(&log_weights_, size * sizeof(float)));
+}
+
+template <>
+__host__ FMBScene<MemoryLocation::HOST>::~FMBScene() {
+    delete[] fmbs_;
+    delete[] log_weights_;
+}
+
+template <>
+__host__ FMBScene<MemoryLocation::DEVICE>::~FMBScene() {
+    CUDA_CHECK(cudaFree(fmbs_));
+    CUDA_CHECK(cudaFree(log_weights_));
+}


The enum type helps us specialize the memory management of the scene type using constructors that are specialized for each MemoryLocation value.

mugamma · 2025-11-27T03:01:46Z

genmetaballs/src/cuda/core/fmb.cuh

@@ -1,20 +1,119 @@
 #pragma once

+#include <cuda/std/span>


I should get rid of this.

mugamma · 2025-11-27T03:04:59Z

genmetaballs/src/cuda/core/fmb.cuh

+        if (x_extent <= 0 || y_extent <= 0 || z_extent <= 0)
+            throw std::domain_error("a metaball cannot have negative extent");
+        extent_ = {x_extent, y_extent, z_extent};


as we move towards the backward pass this will be unnecessary as the FMB will use an inverse-softplus parameterization

Still good to have as a check anyways? In such a case we should never expect this error to trigger. Unless you meant this re: speedup

mugamma · 2025-11-27T03:06:46Z

genmetaballs/src/cuda/core/forward.cu

        float w0 = 0.0f, tf = 0.0f, sumexpd = 0.0f;
        for (const auto& fmb : fmb_getter->get_metaballs(ray)) {
-            const auto& [t, d] = Intersector::intersect(fmb, ray);
+            const auto& [t, d] = Intersector::intersect(fmb, ray, extr);


Note that the intersector now accepts the camera pose.

Nice, so we don't really need to keep the starting location within Ray anymore :)

hooray to memory savings

mugamma · 2025-11-27T03:07:44Z

genmetaballs/src/cuda/core/intersector.cuh

+     */
+    CUDA_CALLABLE static cuda::std::tuple<float, float> intersect(const FMB& fmb, const Ray& ray,
+                                                                  const Pose& cam_pose) {
+        const auto v = cam_pose.get_rot().apply(ray.direction);


This v here is meant to match the in-progress math write-up. I should mention this and fix the rest to match it as well.

Sounds good! We should end up following whatever notation you end up choosing, both in the code and the final write-up

mugamma · 2025-11-27T03:08:47Z

tests/cpp_tests/test_fmb.cu

+    for (auto [fmb, w] : scene) {
+        _num_fmbs += 1;
+    }


I'm just testing here that we can indeed use a range-based for loop over an FMBScene inside a kernel.

horizon-blue

This LGTM! Similar to #24, it'd be great to have some documentations on the intersector :). Feel free to merge this in when you feel ready and I can rebase the camera stuff on top of yours

horizon-blue · 2025-11-28T04:34:35Z

genmetaballs/src/cuda/core/fmb.cu

+CUDA_CALLABLE __forceinline__ Vec3D vecdiv(const Vec3D u, const Vec3D v) {
+    return {u.x / v.x, u.y / v.y, u.z / v.z};
+}


I'm fine with leaving this here since we're not in the header file anyways.

horizon-blue · 2025-11-28T04:40:11Z

genmetaballs/src/cuda/core/forward.cu

        float w0 = 0.0f, tf = 0.0f, sumexpd = 0.0f;
        for (const auto& fmb : fmb_getter->get_metaballs(ray)) {
-            const auto& [t, d] = Intersector::intersect(fmb, ray);
+            const auto& [t, d] = Intersector::intersect(fmb, ray, extr);


Nice, so we don't really need to keep the starting location within Ray anymore :)

horizon-blue · 2025-11-28T04:44:03Z

tests/python_tests/test_intersector.py

+        cov_inv = fmb_rotmat.T @ np.diag(1 / fmb_extent) @ fmb_rotmat
+        t = ((fmb_mu - cam_tran) @ cov_inv @ v) / (v @ cov_inv @ v)
+        d = mahalanobis(fmb_mu, cam_tran + t * v, cov_inv) ** 2
+        # cuda computation


You meant our C++ implementation? (I think this specific instance is actually not running on cuda haha)

arijit-dasgupta · 2025-11-29T05:14:02Z

Sick stuff! iirc the Linear Intersector, is the one used in the paper, so this is good

mugamma commented Nov 27, 2025

View reviewed changes

mugamma requested review from arijit-dasgupta and horizon-blue November 27, 2025 03:11

mugamma marked this pull request as ready for review November 27, 2025 03:11

horizon-blue approved these changes Nov 28, 2025

View reviewed changes

arijit-dasgupta approved these changes Nov 29, 2025

View reviewed changes

mugamma added 3 commits December 2, 2025 18:39

initial version

cd23a82

bindings working

eac6c13

all tests passing

4b6e3a1

horizon-blue force-pushed the mugamma/MET-33-intersector branch from 5c46437 to 4b6e3a1 Compare December 2, 2025 23:40

horizon-blue merged commit f6b7866 into master Dec 2, 2025
1 check passed

horizon-blue deleted the mugamma/MET-33-intersector branch December 2, 2025 23:52

		.def("cov_inv_apply", &FMB::cov_inv_apply,
		"apply the inverse covariance matrix to the given vector", nb::arg("vec"))

		auto [t, d] = LinearIntersector::intersect(fmb, ray, cam_pose);
		return std::make_tuple(t, d);

Conversation

mugamma commented Nov 26, 2025

Uh oh!

linear bot commented Nov 26, 2025

Uh oh!

mugamma left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

horizon-blue left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

horizon-blue Nov 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

arijit-dasgupta commented Nov 29, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

horizon-blue Nov 28, 2025 •

edited

Loading