Adding chunked_pack to pack table chunks when device memory is insufficient by nirandaperera · Pull Request #726 · rapidsai/rapidsmpf

nirandaperera · 2025-12-09T22:07:58Z

This PR enables using cudf::chunked_pack when copying a TableChunk to Host memory, IF there is not enough device memory for a cudf::pack operation.

Signed-off-by: niranda perera <niranda.perera@gmail.com>

copy-pr-bot · 2025-12-09T22:08:02Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…unked-pack

nirandaperera · 2025-12-10T00:10:53Z

cpp/src/integrations/cudf/partition.cpp

+    // TODO: there is a possibility that bounce buffer destructor is a called before the
+    // async copies are completed. Should we synchronize the stream here?
+    bounce_buf->unlock();


@madsbk WDYT?

The operations here are stream ordered. As long as the stream used by table is the same as the stream used by the bounce buffer, everything is correct.

Signed-off-by: niranda perera <niranda.perera@gmail.com>

madsbk

Thanks @nirandaperera, overall looks good.

cpp/src/integrations/cudf/partition.cpp

cpp/include/rapidsmpf/integrations/cudf/partition.hpp

cpp/src/integrations/cudf/partition.cpp

cpp/include/rapidsmpf/integrations/cudf/partition.hpp

cpp/tests/test_partition.cpp

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…unked-pack

nirandaperera · 2025-12-10T18:13:38Z

@madsbk Can you take another look 🙂

cpp/src/integrations/cudf/partition.cpp

madsbk · 2025-12-11T08:12:35Z

cpp/src/streaming/cudf/table_chunk.cpp

+                    if (overbooking > 0) {
+                        // there is not enough memory to pack the table.
+                        size_t avail_dev_mem = pack_res.size() - overbooking;
+                        RAPIDSMPF_EXPECTS(
+                            avail_dev_mem > 1 << 20,
+                            "not enough device memory for the bounce buffer",
+                            std::runtime_error
+                        );
+                        auto bounce_buf = br->allocate(avail_dev_mem, stream(), pack_res);
+
+                        packed_data = std::make_unique<PackedData>(
+                            chunked_pack(table_view(), *bounce_buf, reservation)
+                        );
+                    } else {
+                        // if there is enough memory to pack the table, use `cudf::pack`
+                        auto packed_columns =
+                            cudf::pack(table_view(), stream(), br->device_mr());
+                        // clear the reservation as we are done with it.
+                        pack_res.clear();
+                        packed_data = std::make_unique<PackedData>(
+                            std::move(packed_columns.metadata),
+                            br->move(std::move(packed_columns.gpu_data), stream())
+                        );
+
+                        // Handle the case where `cudf::pack` allocates slightly more than
+                        // the input size. This can occur because cudf uses aligned
+                        // allocations, which may exceed the requested size. To
+                        // accommodate this, we allow some wiggle room.
+                        if (packed_data->data->size > reservation.size()) {
+                            if (packed_data->data->size
+                                <= reservation.size()
+                                       + total_packing_wiggle_room(table_view()))
+                            {
+                                reservation =
+                                    br->reserve(
+                                          MemoryType::HOST, packed_data->data->size, true
+                                    )
+                                        .first;
+                            }


Can we move this to a helper function, I think nesting becomes a problem.
It would also be worth exploring if we even need the non-chunked version?

Yeah, unless there are real performance issues with chunked_pack we should try and always use that. Can we check?

Also agree, let's make a helper function that encapsulates this (so that we can use it outside of tablechunk-copy as well -> this pattern of cudf::pack happens in many places.

Okay, let me add a benchmark.

cpp/src/streaming/cudf/table_chunk.cpp

wence- · 2025-12-11T11:25:27Z

cpp/include/rapidsmpf/integrations/cudf/partition.hpp

+ * @return The total amount of extra memory to reserve for packing.
+ */
+inline size_t total_packing_wiggle_room(cudf::table_view const& table) {
+    return packing_wiggle_room_per_column * static_cast<size_t>(table.num_columns());


note: This is likely not enough if the table has many nested columns.

Sure. I was just mimicking the current impl we have here
https://github.com/rapidsai/rapidsmpf/blob/main/cpp/src/streaming/cudf/table_chunk.cpp#L187-L189

cpp/src/integrations/cudf/partition.cpp

wence- · 2025-12-11T11:29:11Z

cpp/src/integrations/cudf/partition.cpp

+        // all copies are done on the same stream, so we can omit the stream parameter
+        cudf::device_span<uint8_t> buf_span(
+            reinterpret_cast<uint8_t*>(bounce_buf_ptr), chunk_size


Let use use cuda::std::span.

Also the comment is meaningless, device_span has no stream parameter to its ctor.

The comment refers to the stream_view arg in L299 😇

wence- · 2025-12-11T11:33:16Z

cpp/src/streaming/cudf/table_chunk.cpp

-                                    .first;
+                    if (overbooking > 0) {
+                        // there is not enough memory to pack the table.
+                        size_t avail_dev_mem = pack_res.size() - overbooking;


issue: Integer overflow. suppose we can't make a reservation, so pack_res.size() is zero, and overbooking is estimated_memory_usage(...). Then this will be (size_t)(-overbooking).

@wence- I'm not sure if this is true. Reservation is made with overbooking
https://github.com/nirandaperera/rapidsmpf/blob/Make-unbounded-fanout-state-spillable/cpp/src/memory/buffer_resource.cpp#L58-L64
So, it will be

{MemoryReservation(mem_type, this, size), overbooking};

I think the case you are referring to, is for without overbooking

Overbooking is positive. Can it happen that overbooking is larger than pack_res.size()? If yes, then you have integer overflow.

wence- · 2025-12-11T11:36:19Z

cpp/src/streaming/cudf/table_chunk.cpp

+                    if (overbooking > 0) {
+                        // there is not enough memory to pack the table.
+                        size_t avail_dev_mem = pack_res.size() - overbooking;
+                        RAPIDSMPF_EXPECTS(
+                            avail_dev_mem > 1 << 20,
+                            "not enough device memory for the bounce buffer",
+                            std::runtime_error
+                        );
+                        auto bounce_buf = br->allocate(avail_dev_mem, stream(), pack_res);
+
+                        packed_data = std::make_unique<PackedData>(
+                            chunked_pack(table_view(), *bounce_buf, reservation)
+                        );
+                    } else {
+                        // if there is enough memory to pack the table, use `cudf::pack`
+                        auto packed_columns =
+                            cudf::pack(table_view(), stream(), br->device_mr());
+                        // clear the reservation as we are done with it.
+                        pack_res.clear();
+                        packed_data = std::make_unique<PackedData>(
+                            std::move(packed_columns.metadata),
+                            br->move(std::move(packed_columns.gpu_data), stream())
+                        );
+
+                        // Handle the case where `cudf::pack` allocates slightly more than
+                        // the input size. This can occur because cudf uses aligned
+                        // allocations, which may exceed the requested size. To
+                        // accommodate this, we allow some wiggle room.
+                        if (packed_data->data->size > reservation.size()) {
+                            if (packed_data->data->size
+                                <= reservation.size()
+                                       + total_packing_wiggle_room(table_view()))
+                            {
+                                reservation =
+                                    br->reserve(
+                                          MemoryType::HOST, packed_data->data->size, true
+                                    )
+                                        .first;
+                            }


Yeah, unless there are real performance issues with chunked_pack we should try and always use that. Can we check?

Also agree, let's make a helper function that encapsulates this (so that we can use it outside of tablechunk-copy as well -> this pattern of cudf::pack happens in many places.

cpp/tests/test_partition.cpp

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com> Co-authored-by: Lawrence Mitchell <wence@gmx.li>

…unked-pack

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…unked-pack

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…unked-pack

Signed-off-by: niranda perera <niranda.perera@gmail.com>

adding chunked_pack

1d28b4e

Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Dec 9, 2025

nirandaperera added 2 commits December 9, 2025 16:08

adding chunked_pack to table chunk

f6224a8

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

eb8f742

…unked-pack

nirandaperera marked this pull request as ready for review December 10, 2025 00:08

nirandaperera requested a review from a team as a code owner December 10, 2025 00:08

nirandaperera commented Dec 10, 2025

View reviewed changes

nirandaperera added 2 commits December 9, 2025 16:42

fix tests

3f67fce

Signed-off-by: niranda perera <niranda.perera@gmail.com>

minor change

241d697

Signed-off-by: niranda perera <niranda.perera@gmail.com>

madsbk requested changes Dec 10, 2025

View reviewed changes

nirandaperera and others added 3 commits December 10, 2025 09:39

Apply suggestions from code review

d12c5eb

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com>

addressing PR comments

2a08a1d

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

88fa3e4

…unked-pack

nirandaperera requested a review from madsbk December 10, 2025 18:12

madsbk reviewed Dec 11, 2025

View reviewed changes

madsbk mentioned this pull request Dec 11, 2025

Introducing MemoryType::PINNED_HOST #731

Merged

wence- requested changes Dec 11, 2025

View reviewed changes

Apply suggestions from code review

bf5088d

Co-authored-by: Mads R. B. Kristensen <madsbk@gmail.com> Co-authored-by: Lawrence Mitchell <wence@gmx.li>

nirandaperera requested review from madsbk and wence- December 11, 2025 18:20

nirandaperera added 7 commits December 11, 2025 14:56

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

027875e

…unked-pack

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

856f515

…unked-pack

adding pack_to_host

5f984f6

Signed-off-by: niranda perera <niranda.perera@gmail.com>

finalize packing

78ed513

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

34c96cc

…unked-pack

Merge branch 'main' into use-cudf-chunked-pack

41760de

Merge branch 'main' into use-cudf-chunked-pack

a897597

nirandaperera added 3 commits January 7, 2026 16:36

fix build

d35feb5

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'main' of github.com:rapidsai/rapidsmpf into use-cudf-ch…

8aa081d

…unked-pack

fixing tests

e1fcf23

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Conversation

nirandaperera commented Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot bot commented Dec 9, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

madsbk left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nirandaperera commented Dec 10, 2025

Uh oh!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nirandaperera commented Dec 9, 2025 •

edited

Loading