Adding launch params as argumetns to forward

Author: horizon-blue

genmetaballs/src/cuda/bindings.cu

@@ -228,6 +228,17 @@ NB_MODULE(_genmetaballs_bindings, m) {
     bind_array2d<float, MemoryLocation::HOST>(utils, "CPUFloatArray2D");
     bind_array2d<float, MemoryLocation::DEVICE>(utils, "GPUFloatArray2D");
 
+    // bind dim3, which is used to specify the launch configuration for the kernel
+    nb::class_<dim3>(utils, "dim3")
+        .def(nb::init<uint32_t, uint32_t, uint32_t>(), nb::arg("x") = 1, nb::arg("y") = 1,
+             nb::arg("z") = 1)
+        .def_prop_ro("x", [](const dim3& self) { return self.x; })
+        .def_prop_ro("y", [](const dim3& self) { return self.y; })
+        .def_prop_ro("z", [](const dim3& self) { return self.z; })
+        .def("__repr__", [](const dim3& self) {
+            return nb::str("dim3(x={}, y={}, z={})").format(self.x, self.y, self.z);
+        });
+
 } // NB_MODULE(_genmetaballs_bindings)
 
 template <typename T, MemoryLocation location>
@@ -319,5 +330,5 @@ void bind_render_fmbs(nb::module_& m, const char* name) {
           &render_fmbs<AllGetter<MemoryLocation::DEVICE>, LinearIntersector, Blender, Confidence>,
           "Render the given FMB scene into the provided image view", nb::arg("fmbs"),
           nb::arg("blender"), nb::arg("confidence"), nb::arg("intr"), nb::arg("extr"),
-          nb::arg("img"));
+          nb::arg("img"), nb::arg("grid_size"), nb::arg("block_size"));
 }

genmetaballs/src/cuda/core/forward.cuh

@@ -9,9 +9,6 @@
 #include "image.cuh"
 #include "utils.cuh"
 
-// TODO: tune this number
-constexpr auto NUM_BLOCKS = dim3(4, 4);
-constexpr auto THREADS_PER_BLOCK = dim3(16, 16);
 
 CUDA_CALLABLE PixelCoordRange get_pixel_coords(const dim3 thread_idx, const dim3 block_idx,
                                                const dim3 block_dim, const dim3 grid_dim,
@@ -49,7 +46,8 @@ __global__ void render_kernel(const FMBScene<MemoryLocation::DEVICE>& fmbs, cons
 template <typename Getter, typename Intersector, typename Blender, typename Confidence>
 void render_fmbs(const FMBScene<MemoryLocation::DEVICE>& fmbs, const Blender& blender,
                  const Confidence& confidence, const Intrinsics& intr, const Pose& extr,
-                 ImageView<MemoryLocation::DEVICE> img) {
+                 ImageView<MemoryLocation::DEVICE> img, const dim3 grid_size,
+                 const dim3 block_size) {
     render_kernel<Getter, Intersector, Blender, Confidence>
-        <<<NUM_BLOCKS, THREADS_PER_BLOCK>>>(fmbs, blender, confidence, intr, extr, img);
+        <<<grid_size, block_size>>>(fmbs, blender, confidence, intr, extr, img);
 }

genmetaballs/src/genmetaballs/core/__init__.py

@@ -12,7 +12,12 @@
 )
 from genmetaballs._genmetaballs_bindings.fmb import FMB, CPUFMBScene, GPUFMBScene
 from genmetaballs._genmetaballs_bindings.image import CPUImage, GPUImage
-from genmetaballs._genmetaballs_bindings.utils import CPUFloatArray2D, GPUFloatArray2D, sigmoid
+from genmetaballs._genmetaballs_bindings.utils import (
+    CPUFloatArray2D,
+    GPUFloatArray2D,
+    dim3,
+    sigmoid,
+)
 
 type DeviceType = Literal["cpu", "gpu"]
 
@@ -83,6 +88,8 @@ def render_fmbs(
     intr: Intrinsics,
     extr: geometry.Pose,
     img: GPUImage | None = None,
+    grid_size: dim3 = dim3(4, 4),
+    block_size: dim3 = dim3(16, 16),
 ) -> GPUImage:
     """Render the given FMB scene into the provided image view.
 
@@ -105,7 +112,7 @@ def render_fmbs(
     else:
         raise TypeError("Unsupported blender and confidence combination.")
 
-    render_func(fmbs, blender, confidence, intr, extr, img.as_view())
+    render_func(fmbs, blender, confidence, intr, extr, img.as_view(), grid_size, block_size)
     return img
 
 
@@ -122,6 +129,7 @@ def render_fmbs(
     "Camera",
     "FourParameterBlender",
     "FMB",
+    "dim3",
     "Intrinsics",
     "ThreeParameterBlender",
     "TwoParameterConfidence",