diff --git a/CMakeLists.txt b/CMakeLists.txt index 054ac28..73ee145 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -19,6 +19,7 @@ add_library(genmetaballs_core genmetaballs/src/cuda/core/add.cuh genmetaballs/src/cuda/core/geometry.cuh genmetaballs/src/cuda/core/geometry.cu + genmetaballs/src/cuda/core/confidence.cuh ) # Set include directories for the core library diff --git a/genmetaballs/src/cuda/bindings.cu b/genmetaballs/src/cuda/bindings.cu index 15cbcc0..f75a322 100644 --- a/genmetaballs/src/cuda/bindings.cu +++ b/genmetaballs/src/cuda/bindings.cu @@ -2,9 +2,12 @@ #include #include #include +#include #include "core/add.cuh" +#include "core/confidence.cuh" #include "core/geometry.cuh" +#include "core/utils.cuh" constexpr uint32_t GRID_DIM = 4096; constexpr uint32_t BLOCK_DIM = 1024; @@ -12,9 +15,12 @@ constexpr uint32_t BLOCK_DIM = 1024; namespace nb = nanobind; NB_MODULE(_genmetaballs_bindings, m) { + + // simple add kernel m.def("gpu_add", &gpu_add, "Add two lists elementwise on the GPU", nb::arg("a"), nb::arg("b")); + // exposing Vec3D nb::class_(m, "Vec3D") .def(nb::init<>()) .def(nb::init()) @@ -23,8 +29,21 @@ NB_MODULE(_genmetaballs_bindings, m) { .def_rw("z", &Vec3D::z) .def(nb::self + nb::self) .def(nb::self - nb::self) - .def("__repr__", [](const Vec3D& v) { - nb::str s = nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); - return s; - }); -} + .def("__repr__", + [](const Vec3D& v) { return nb::str("Vec3D({}, {}, {})").format(v.x, v.y, v.z); }); + + // confidence submodule + nb::module_ confidence = m.def_submodule("confidence"); + nb::class_(confidence, "ZeroParameterConfidence") + .def(nb::init<>()) + .def("get_confidence", &ZeroParameterConfidence::get_confidence); + + nb::class_(confidence, "TwoParameterConfidence") + .def(nb::init()) + .def("get_confidence", &TwoParameterConfidence::get_confidence); + + // utils submodule + nb::module_ utils = m.def_submodule("utils"); + utils.def("sigmoid", sigmoid, nb::arg("x"), "Compute the sigmoid function: 1 / (1 + exp(-x))"); + +} // NB_MODULE(_genmetaballs_bindings) diff --git a/genmetaballs/src/cuda/core/blender.cuh b/genmetaballs/src/cuda/core/blender.cuh index 8c4e21f..5ce8c7a 100644 --- a/genmetaballs/src/cuda/core/blender.cuh +++ b/genmetaballs/src/cuda/core/blender.cuh @@ -8,7 +8,7 @@ struct ThreeParameterBlender { float beta2; float eta; - __host__ __device__ __forceinline__ // TODO inline? + CUDA_CALLABLE __forceinline__ // TODO inline? float blend(float t, float d, const FMB& fmb, const Ray& ray) const; }; diff --git a/genmetaballs/src/cuda/core/confidence.cuh b/genmetaballs/src/cuda/core/confidence.cuh index b79c3ad..5320b0e 100644 --- a/genmetaballs/src/cuda/core/confidence.cuh +++ b/genmetaballs/src/cuda/core/confidence.cuh @@ -1,12 +1,23 @@ #pragma once #include +#include +#include + +#include "utils.cuh" struct TwoParameterConfidence { + float beta4; float beta5; + CUDA_CALLABLE __forceinline__ float get_confidence(float sumexpd) const { + return sigmoid(beta4 * sumexpd + beta5); + } +}; + +struct ZeroParameterConfidence { - __host__ __device__ __forceinline__ float get_confidence(float sumexpd) { - return 0; - } // TODO + CUDA_CALLABLE __forceinline__ float get_confidence(float sumexpd) const { + return 1.0f - expf(-sumexpd); + } }; diff --git a/genmetaballs/src/cuda/core/forward.cu b/genmetaballs/src/cuda/core/forward.cu index e2b9117..a71caf6 100644 --- a/genmetaballs/src/cuda/core/forward.cu +++ b/genmetaballs/src/cuda/core/forward.cu @@ -1,12 +1,13 @@ #include #include +#include constexpr NUM_BLOCKS dim3(10); // XXX madeup constexpr THREADS_PER_BLOCK dim3(10); namespace FMB { -__device__ __host__ std::vector> get_pixel_coords_and_rays( +CUDA_CALLABLE std::vector> get_pixel_coords_and_rays( const dim3 thread_idx, const dim3 block_idx) { std::vector> res; diff --git a/genmetaballs/src/cuda/core/intersector.cuh b/genmetaballs/src/cuda/core/intersector.cuh index acfb3a1..8d71b66 100644 --- a/genmetaballs/src/cuda/core/intersector.cuh +++ b/genmetaballs/src/cuda/core/intersector.cuh @@ -8,6 +8,5 @@ // implement equation (6) in the paper class LinearIntersector { - static __device__ __host__ std::pair intersect(const FMB& fmb, - const Ray& ray) const; + static CUDA_CALLABLE std::pair intersect(const FMB& fmb, const Ray& ray) const; }; diff --git a/genmetaballs/src/cuda/core/utils.cuh b/genmetaballs/src/cuda/core/utils.cuh index 108cd4d..1dcdf36 100644 --- a/genmetaballs/src/cuda/core/utils.cuh +++ b/genmetaballs/src/cuda/core/utils.cuh @@ -1,9 +1,12 @@ #pragma once +#include #include #include #include +#define CUDA_CALLABLE __host__ __device__ + #define CUDA_CHECK(x) \ do { \ cuda_check((x), __FILE__, __LINE__); \ @@ -11,6 +14,13 @@ void cuda_check(cudaError_t code, const char* file, int line); +CUDA_CALLABLE __forceinline__ float sigmoid(float x) { + if (isnan(x)) { + return x; + } + return 1.0f / (1.0f + expf(-x)); +} + // Non-owning 2D view into a contiguous array in either host or device memory template class Array2D { diff --git a/genmetaballs/src/genmetaballs/core/__init__.py b/genmetaballs/src/genmetaballs/core/__init__.py new file mode 100644 index 0000000..341a681 --- /dev/null +++ b/genmetaballs/src/genmetaballs/core/__init__.py @@ -0,0 +1,11 @@ +from genmetaballs._genmetaballs_bindings.confidence import ( + TwoParameterConfidence, + ZeroParameterConfidence, +) +from genmetaballs._genmetaballs_bindings.utils import sigmoid + +__all__ = [ + "ZeroParameterConfidence", + "TwoParameterConfidence", + "sigmoid", +] diff --git a/tests/cpp_tests/test_confidence.cu b/tests/cpp_tests/test_confidence.cu new file mode 100644 index 0000000..cda499a --- /dev/null +++ b/tests/cpp_tests/test_confidence.cu @@ -0,0 +1,147 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/confidence.cuh" + +// Helper: Python ground truth, as in test_confidence.py +inline float ground_truth_expit(float x) { + return 1.0F / (1.0F + std::exp(-x)); +} +float ground_truth_two_parameter_confidence(float beta4, float beta5, float sumexpd) { + return ground_truth_expit((beta4 * sumexpd) + beta5); +} +float ground_truth_zero_parameter_confidence(float sumexpd) { + return 1.0F - std::exp(-sumexpd); +} + +template +__global__ void confidence_kernel(const float* sumexpd, float* confidences, uint32_t n, + Confidence confidence) { + uint32_t i = threadIdx.x + (blockIdx.x * blockDim.x); + if (i < n) { + confidences[i] = confidence.get_confidence(sumexpd[i]); + } +} + +constexpr uint32_t GRID_DIM = 256; +constexpr uint32_t BLOCK_DIM = 1024; + +template +std::vector gpu_get_confidence(const std::vector& sumexpd_vec, + Confidence confidence) { + auto n = static_cast(sumexpd_vec.size()); + auto nbytes = n * sizeof(float); + float *d_sumexpd = nullptr, *d_confidences = nullptr; + std::vector result(n); + + CUDA_CHECK(cudaMalloc(&d_sumexpd, nbytes)); + CUDA_CHECK(cudaMalloc(&d_confidences, nbytes)); + CUDA_CHECK(cudaMemcpy(d_sumexpd, sumexpd_vec.data(), nbytes, cudaMemcpyHostToDevice)); + + auto block_dim = BLOCK_DIM; + auto grid_dim = (n + block_dim - 1) / block_dim; + if (grid_dim > GRID_DIM) + grid_dim = GRID_DIM; + + confidence_kernel<<>>(d_sumexpd, d_confidences, n, confidence); + + CUDA_CHECK(cudaMemcpy(result.data(), d_confidences, nbytes, cudaMemcpyDeviceToHost)); + CUDA_CHECK(cudaFree(d_sumexpd)); + CUDA_CHECK(cudaFree(d_confidences)); + return result; +} + +constexpr int NUM_RNG_SEEDS_PER_TEST = 5; +constexpr int NUM_N_VALUES_PER_TEST = 5; +constexpr uint32_t MASTER_SEED = 0; + +static std::vector confidence_test_sizes() { + std::vector sizes; + for (int k = 0; k < NUM_N_VALUES_PER_TEST; ++k) + sizes.push_back(1 << (4 + k)); // 2^(4+k): [16, 32, 64, 128, 256] + return sizes; +} + +// Define simple struct to match python CONFIDENCE_TEST_CASES +struct ConfidenceCase { + std::string name; + float beta4 = 0.0F; + float beta5 = 0.0F; + bool is_two_param; +}; +static std::vector confidence_cases() { + return { + {"zero_param", 0.0F, 0.0F, false}, + {"two_param_0.5_-1", 0.5F, -1.0F, true}, + {"two_param_1_0", 1.0F, 0.0F, true}, + {"two_param_-0.5_2", -0.5F, 2.0F, true}, + }; +} + +TEST(GpuConfidenceTest, ConfidenceMultipleValuesGPU_AllTypes) { + using test_float = float; + constexpr float rtol = 1e-6F; + + auto sizes = confidence_test_sizes(); + std::mt19937 master_gen(MASTER_SEED); + std::uniform_int_distribution seed_dist(0, std::numeric_limits::max()); + std::vector seeds(NUM_RNG_SEEDS_PER_TEST); + for (auto& s : seeds) + s = seed_dist(master_gen); + + for (int size_idx = 0; size_idx < static_cast(sizes.size()); ++size_idx) { + int N = sizes[size_idx]; + + for (const auto& conf_case : confidence_cases()) { + for (uint32_t test_seed : seeds) { + SCOPED_TRACE(testing::Message() << "N=" << N << ", seed=" << test_seed + << ", conf_type=" << conf_case.name); + + // Use float32 min as lower and float32 max as upper bound + std::mt19937 rng(test_seed); + std::uniform_real_distribution dist(std::numeric_limits::min(), + std::numeric_limits::max()); + std::vector sumexpd_vec(N); + for (int i = 0; i < N; ++i) + sumexpd_vec[i] = dist(rng); + + std::vector expected(N); + if (conf_case.is_two_param) { + for (int i = 0; i < N; ++i) + expected[i] = ground_truth_two_parameter_confidence( + conf_case.beta4, conf_case.beta5, sumexpd_vec[i]); + } else { + for (int i = 0; i < N; ++i) + expected[i] = ground_truth_zero_parameter_confidence(sumexpd_vec[i]); + } + + std::vector actual; + if (conf_case.is_two_param) { + TwoParameterConfidence conf(conf_case.beta4, conf_case.beta5); + actual = gpu_get_confidence(sumexpd_vec, conf); + } else { + ZeroParameterConfidence conf; + actual = gpu_get_confidence(sumexpd_vec, conf); + } + + ASSERT_EQ(actual.size(), expected.size()); + for (int i = 0; i < N; ++i) { + ASSERT_NEAR(actual[i], expected[i], 1e-6F) + << "at idx=" << i << " N=" << N << " conf_type=" << conf_case.name + << " exp=" << expected[i] << " act=" << actual[i]; + } + // Ensure all actual values are in [0, 1] + ASSERT_TRUE(std::all_of(actual.begin(), actual.end(), + [](float v) { return v >= 0.0F && v <= 1.0F; })) + << "out-of-bounds value(s) detected for conf_type=" << conf_case.name; + } + } + } +} diff --git a/tests/cpp_tests/test_utils.cu b/tests/cpp_tests/test_utils.cu index 9fb97e8..575502a 100644 --- a/tests/cpp_tests/test_utils.cu +++ b/tests/cpp_tests/test_utils.cu @@ -1,6 +1,9 @@ +#include +#include #include #include #include +#include #include #include #include @@ -8,6 +11,105 @@ #include "core/utils.cuh" +namespace test_utils_gpu { + +// CUDA kernel for computing sigmoid element-wise (relies on __device__ sigmoid in utils.cuh) +__global__ void sigmoid_kernel(const float* x, float* result, uint32_t n) { + uint32_t i = threadIdx.x + blockIdx.x * blockDim.x; + if (i < n) { + result[i] = sigmoid(x[i]); + } +} + +// GPU function to compute sigmoid for a vector (float only) +template +std::vector gpu_sigmoid(const std::vector& x_vec) { + uint32_t n = x_vec.size(); + uint32_t nbytes = n * sizeof(float); + float *d_x = nullptr, *d_result = nullptr; + std::vector result(n); + + CUDA_CHECK(cudaMalloc(&d_x, nbytes)); + CUDA_CHECK(cudaMalloc(&d_result, nbytes)); + + CUDA_CHECK(cudaMemcpy(d_x, x_vec.data(), nbytes, cudaMemcpyHostToDevice)); + + sigmoid_kernel<<>>(d_x, d_result, n); + CUDA_CHECK(cudaDeviceSynchronize()); + + CUDA_CHECK(cudaMemcpy(result.data(), d_result, nbytes, cudaMemcpyDeviceToHost)); + + CUDA_CHECK(cudaFree(d_x)); + CUDA_CHECK(cudaFree(d_result)); + + return result; +} + +// Host sigmoid for reference +inline float host_sigmoid(float x) { + return 1.0f / (1.0f + std::exp(-x)); +} + +} // namespace test_utils_gpu + +// Parameters matching the removed Python test +constexpr int NUM_RNG_SEEDS_PER_TEST = 5; +constexpr int NUM_N_VALUES_PER_TEST = 5; +constexpr int SEED_MASTER = 0; + +// Helper: Generate test sizes [16, 32, 64, 128, 256] +static std::vector sigmoid_test_sizes() { + std::vector sizes; + for (int k = 0; k < NUM_N_VALUES_PER_TEST; ++k) + sizes.push_back(1 << (4 + k)); // 2^(4+k) + return sizes; +} + +TEST(GpuSigmoidTest, SigmoidVectorCorrectness) { + // Generate seeds + std::mt19937 master_gen(SEED_MASTER); + std::uniform_int_distribution seed_dist(0, std::numeric_limits::max()); + std::vector seeds(NUM_RNG_SEEDS_PER_TEST); + for (auto& s : seeds) + s = seed_dist(master_gen); + + auto sizes = sigmoid_test_sizes(); + + for (size_t size_idx = 0; size_idx < sizes.size(); ++size_idx) { + int N = sizes[size_idx]; + for (uint32_t seed : seeds) { + // Create reproducible random numbers in [-10, 10] + std::mt19937 rng(seed); + std::uniform_real_distribution dist(-10.0f, 10.0f); + std::vector x_vec(N); + for (int i = 0; i < N; ++i) + x_vec[i] = dist(rng); + + // Compute expected (host) + std::vector expected(N); + for (int i = 0; i < N; ++i) + expected[i] = test_utils_gpu::host_sigmoid(x_vec[i]); + + // Compute actual (GPU) + constexpr uint32_t block_dim = 256; + uint32_t grid_dim = (N + block_dim - 1) / block_dim; + std::vector actual = test_utils_gpu::gpu_sigmoid<1024, block_dim>(x_vec); + + // Compare + ASSERT_EQ(actual.size(), expected.size()); + for (int i = 0; i < N; ++i) { + ASSERT_NEAR(actual[i], expected[i], 1e-5) + << "at idx=" << i << " for N=" << N << " seed=" << seed; + } + + // Check [0, 1] bounds + ASSERT_TRUE(std::all_of(actual.begin(), actual.end(), + [](float v) { return v >= 0.0f && v <= 1.0f; })) + << "Sigmoid output out of [0,1] range for N=" << N << " seed=" << seed; + } + } +} + // CUDA kernel to fill Array2D with sequential values __global__ void fill_array2d_kernel(Array2D array2d) { uint32_t i = threadIdx.x; diff --git a/tests/python_tests/test_confidence.py b/tests/python_tests/test_confidence.py new file mode 100644 index 0000000..4ccd88a --- /dev/null +++ b/tests/python_tests/test_confidence.py @@ -0,0 +1,96 @@ +import numpy as np +import pytest +from scipy.special import expit + +from genmetaballs.core import ( + TwoParameterConfidence, + ZeroParameterConfidence, +) + + +def ground_truth_two_parameter_confidence( + beta4: float, beta5: float, sumexpd: float | np.ndarray +) -> float | np.ndarray: + """Compute the five parameter confidence for a single value or an array of values using numpy.""" + return expit((beta4 * sumexpd) + beta5) + + +def ground_truth_zero_parameter_confidence( + sumexpd: float | np.ndarray, +) -> float | np.ndarray: + """Compute the three parameter confidence for a single value or an array of values using numpy.""" + return 1.0 - np.exp(-sumexpd) + + +NUM_RNG_SEEDS_PER_TEST = 5 +NUM_N_VALUES_PER_TEST = 5 +MASTER_SEED = 0 + + +# Test data for different confidence types +CONFIDENCE_TEST_CASES = [ + # (confidence_class, confidence_kwargs, ground_truth_func, ground_truth_kwargs) + ("zero_param", {}, ground_truth_zero_parameter_confidence, {}), + ( + "two_param", + {"beta4": 0.5, "beta5": -1.0}, + ground_truth_two_parameter_confidence, + {"beta4": 0.5, "beta5": -1.0}, + ), + ( + "two_param", + {"beta4": 1.0, "beta5": 0.0}, + ground_truth_two_parameter_confidence, + {"beta4": 1.0, "beta5": 0.0}, + ), + ( + "two_param", + {"beta4": -0.5, "beta5": 2.0}, + ground_truth_two_parameter_confidence, + {"beta4": -0.5, "beta5": 2.0}, + ), +] + + +def create_confidence_instance(conf_type: str, kwargs: dict): + """Helper function to dispatch the appropriate confidence instance.""" + if conf_type == "two_param": + return TwoParameterConfidence(kwargs["beta4"], kwargs["beta5"]) + elif conf_type == "zero_param": + return ZeroParameterConfidence() + else: + raise ValueError(f"Unknown confidence type: {conf_type}") + + +@pytest.mark.parametrize( + "rng_seed", np.random.default_rng(MASTER_SEED).integers(0, 2**32, size=NUM_RNG_SEEDS_PER_TEST) +) +@pytest.mark.parametrize( + "conf_type, conf_kwargs, ground_truth_func, gt_kwargs", CONFIDENCE_TEST_CASES +) +def test_confidence_single_value( + rng_seed: int, conf_type: str, conf_kwargs: dict, ground_truth_func, gt_kwargs: dict +) -> None: + """Test that confidence can be computed correctly on the CPU for a single value across all confidence types.""" + rng = np.random.default_rng(rng_seed) + confidence = create_confidence_instance(conf_type, conf_kwargs) + + # range of sumexpd is [tiny (smallest f32 value above 0), max (largest f32 value)] since it is the sum of exp(d) for all metaballs + sumexpd = ( + rng.uniform(low=np.finfo(np.float32).tiny, high=np.finfo(np.float32).max, size=1) + .astype(np.float32) + .item() + ) + + # Compute expected using appropriate ground truth function + if conf_type == "two_param": + expected = ground_truth_func(gt_kwargs["beta4"], gt_kwargs["beta5"], sumexpd) + else: + expected = ground_truth_func(sumexpd) + + actual = confidence.get_confidence(sumexpd) + + # check that the actual and expected values are close + assert np.isclose(actual, expected, rtol=1e-6) + # check that all confidence values are between 0 and 1 inclusive + assert actual >= 0.0 and actual <= 1.0 diff --git a/tests/python_tests/test_utils.py b/tests/python_tests/test_utils.py new file mode 100644 index 0000000..41d2714 --- /dev/null +++ b/tests/python_tests/test_utils.py @@ -0,0 +1,60 @@ +import numpy as np +import pytest +from scipy.special import expit + +from genmetaballs.core import sigmoid + +NUM_RNG_SEEDS_PER_TEST = 5 +NUM_N_VALUES_PER_TEST = 5 +MASTER_SEED = 0 + + +@pytest.mark.parametrize( + "rng_seed", np.random.default_rng(MASTER_SEED).integers(0, 2**32, size=NUM_RNG_SEEDS_PER_TEST) +) +def test_sigmoid_single_value(rng_seed: int) -> None: + """Test that sigmoid can be computed correctly for a single value.""" + rng = np.random.default_rng(rng_seed) + # Test with a wide range of values + x = rng.uniform(low=-10.0, high=10.0, size=1).astype(np.float32).item() + + # Compute expected result using scipy + expected = expit(x) + + # Compute actual result using our implementation + actual = sigmoid(x) + + # Compare results - use reasonable tolerance for float32 + assert np.isclose(actual, expected, rtol=1e-5, atol=1e-6) + + # Check that sigmoid output is within [0, 1] + assert actual >= 0.0 + assert actual <= 1.0 + + +@pytest.mark.parametrize( + "x", + [ + -1e30, + -1e10, + -1e-30, + 0.0, + 1e-30, + 1e10, + 1e30, + float("-inf"), + float("inf"), + float("nan"), + ], +) +def test_sigmoid_edge_cases(x: float) -> None: + """Test sigmoid with edge case values.""" + expected = expit(x) + actual = sigmoid(x) + + if np.isnan(expected): + assert np.isnan(actual) + else: + assert np.isclose(actual, expected, rtol=1e-5, atol=1e-6) + assert actual >= 0.0 + assert actual <= 1.0