From df48ef1a191bb4e9336640f6029731010bb0c459 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 28 Jun 2026 11:40:32 +0000 Subject: [PATCH 1/3] Add memory resource control to valid_if Teach cudf::detail::valid_if to accept cudf::memory_resources so the allocation roles inside the helper can be controlled independently. The returned validity mask now uses the output resource, while the device scalar used to accumulate the valid count uses the temporary resource. Replacing the former device_async_resource_ref parameter preserves existing source calls through memory_resources' implicit one-resource construction. Those callers retain the previous behavior: their supplied resource owns the returned mask and the current cuDF resource supplies temporary storage. Explicit two-resource callers can now avoid any fallback to global resource state. Add focused bitmask tests using the reusable memory-resource harness. The non-empty case installs an allocation-failing current resource, verifies output allocation lifetime, confirms temporary allocation activity and release, and validates the generated mask. The empty case proves the explicit two-resource path remains allocation-free and does not query the current resource. Validation performed: - ninja -C cpp/build/latest BITMASK_TEST - cpp/build/latest/gtests/BITMASK_TEST --gtest_filter='ValidIfTest.Explicit*' - test-cudf-cpp -R BITMASK_TEST - pre-commit run --all-files --- cpp/include/cudf/detail/valid_if.cuh | 12 +++--- cpp/tests/bitmask/valid_if_tests.cu | 57 ++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 6 deletions(-) diff --git a/cpp/include/cudf/detail/valid_if.cuh b/cpp/include/cudf/detail/valid_if.cuh index 09f3afefdc1b..94a6e2f6a11e 100644 --- a/cpp/include/cudf/detail/valid_if.cuh +++ b/cpp/include/cudf/detail/valid_if.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -73,7 +73,7 @@ CUDF_KERNEL void valid_if_kernel( * @param end The end of the sequence * @param p The predicate * @param stream CUDA stream used for device memory operations and kernel launches - * @param mr Device memory resource used to allocate the returned device memory + * @param resources Memory resources used for the returned bitmask and temporary count scalar * @return A pair containing a `device_buffer` with the new bitmask and its null count */ template @@ -81,18 +81,18 @@ std::pair valid_if(InputIterator begin, InputIterator end, Predicate p, rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) + cudf::memory_resources resources) { CUDF_EXPECTS(begin <= end, "Invalid range."); size_type size = cuda::std::distance(begin, end); - auto null_mask = cudf::create_null_mask(size, mask_state::UNINITIALIZED, stream, mr); + auto null_mask = + cudf::create_null_mask(size, mask_state::UNINITIALIZED, stream, resources.get_output_mr()); size_type null_count{0}; if (size > 0) { - cudf::detail::device_scalar valid_count{ - 0, stream, cudf::get_current_device_resource_ref()}; + cudf::detail::device_scalar valid_count{0, stream, resources.get_temporary_mr()}; constexpr size_type block_size{256}; grid_1d grid{size, block_size}; diff --git a/cpp/tests/bitmask/valid_if_tests.cu b/cpp/tests/bitmask/valid_if_tests.cu index ffcdc0b4293a..a32fdcdbc9ed 100644 --- a/cpp/tests/bitmask/valid_if_tests.cu +++ b/cpp/tests/bitmask/valid_if_tests.cu @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -42,6 +43,30 @@ TEST_F(ValidIfTest, EmptyRange) EXPECT_EQ(0, actual.second); } +TEST_F(ValidIfTest, ExplicitMemoryResourcesEmptyRange) +{ + auto harness = cudf::test::memory_resource_test_harness{this->mr()}; + auto stream = cudf::get_default_stream(); + + auto actual = [&] { + auto current_scope = harness.fail_on_current_device_resource_use(); + auto result = cudf::detail::valid_if(cuda::counting_iterator{0}, + cuda::counting_iterator{0}, + odds_valid{}, + stream, + harness.resources()); + harness.synchronize(stream); + return result; + }(); + + EXPECT_EQ(0u, actual.first.size()); + EXPECT_EQ(nullptr, actual.first.data()); + EXPECT_EQ(0, actual.second); + harness.expect_no_live_allocations(stream); + EXPECT_EQ(0, harness.output_mr().get_bytes_counter().total); + EXPECT_EQ(0, harness.temporary_mr().get_bytes_counter().total); +} + TEST_F(ValidIfTest, InvalidRange) { EXPECT_THROW(cudf::detail::valid_if(cuda::counting_iterator{1}, @@ -66,6 +91,38 @@ TEST_F(ValidIfTest, OddsValid) EXPECT_EQ(expected.second, actual.second); } +TEST_F(ValidIfTest, ExplicitMemoryResourceControl) +{ + auto harness = cudf::test::memory_resource_test_harness{this->mr()}; + auto stream = cudf::get_default_stream(); + auto comparison_resources = cudf::memory_resources{harness.setup_mr(), harness.setup_mr()}; + auto iter = cudf::detail::make_counting_transform_iterator(0, odds_valid{}); + auto expected = cudf::test::detail::make_null_mask(iter, iter + 10000, comparison_resources); + + { + auto actual = [&] { + auto current_scope = harness.fail_on_current_device_resource_use(); + auto result = cudf::detail::valid_if(cuda::counting_iterator{0}, + cuda::counting_iterator{10000}, + odds_valid{}, + stream, + harness.resources()); + harness.synchronize(stream); + return result; + }(); + + harness.expect_output_allocations_live(stream); + harness.expect_temporary_allocation_activity(stream); + harness.expect_temporary_allocations_released(stream); + CUDF_TEST_EXPECT_EQUAL_BUFFERS( + expected.first.data(), actual.first.data(), expected.first.size(), comparison_resources); + EXPECT_EQ(5000, actual.second); + EXPECT_EQ(expected.second, actual.second); + } + + harness.expect_no_live_allocations(stream); +} + TEST_F(ValidIfTest, AllValid) { auto iter = cudf::detail::make_counting_transform_iterator(0, all_valid{}); From 9b4483bae4f048ec09758139f26ac29a78e28ad3 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 11 Aug 2026 09:39:59 -0500 Subject: [PATCH 2/3] Fix valid_if test copyright notice --- cpp/tests/bitmask/valid_if_tests.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/bitmask/valid_if_tests.cu b/cpp/tests/bitmask/valid_if_tests.cu index a32fdcdbc9ed..d328469e3dc1 100644 --- a/cpp/tests/bitmask/valid_if_tests.cu +++ b/cpp/tests/bitmask/valid_if_tests.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From d94928052bc502ebd890868ba13b3d61a5345a92 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 16 Aug 2026 12:04:53 -0500 Subject: [PATCH 3/3] Fix valid_if test helpers after rebase --- cpp/tests/bitmask/valid_if_tests.cu | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/cpp/tests/bitmask/valid_if_tests.cu b/cpp/tests/bitmask/valid_if_tests.cu index d328469e3dc1..6d6dd4aee422 100644 --- a/cpp/tests/bitmask/valid_if_tests.cu +++ b/cpp/tests/bitmask/valid_if_tests.cu @@ -97,7 +97,8 @@ TEST_F(ValidIfTest, ExplicitMemoryResourceControl) auto stream = cudf::get_default_stream(); auto comparison_resources = cudf::memory_resources{harness.setup_mr(), harness.setup_mr()}; auto iter = cudf::detail::make_counting_transform_iterator(0, odds_valid{}); - auto expected = cudf::test::detail::make_null_mask(iter, iter + 10000, comparison_resources); + auto expected = + cudf::test::detail::make_null_mask(iter, iter + 10000, stream, comparison_resources); { auto actual = [&] { @@ -114,8 +115,11 @@ TEST_F(ValidIfTest, ExplicitMemoryResourceControl) harness.expect_output_allocations_live(stream); harness.expect_temporary_allocation_activity(stream); harness.expect_temporary_allocations_released(stream); - CUDF_TEST_EXPECT_EQUAL_BUFFERS( - expected.first.data(), actual.first.data(), expected.first.size(), comparison_resources); + CUDF_TEST_EXPECT_EQUAL_BUFFERS(expected.first.data(), + actual.first.data(), + expected.first.size(), + stream, + comparison_resources); EXPECT_EQ(5000, actual.second); EXPECT_EQ(expected.second, actual.second); }