NVIDIA / NVIDIA/cudf

[FEA] Simplified packing API and impl

Open
#21,874 2 comments 0 reactions 1 assignee Claimed by @nirandaperera View on GitHub
feature request libcudf no-oom
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**
**Describe the solution you'd like**

# Current API

The `cudf::chunked_pack` class is designed for performing chunked packing operations on a `table_view` into a user-provided buffer, primarily for managing data movement under GPU memory constraints or to mitigate out-of-memory errors.
[https://github.com/rapidsai/cudf/blob/main/cpp/include/cudf/contiguous\_split.hpp\#L152-L228](https://github.com/rapidsai/cudf/blob/main/cpp/include/cudf/contiguous_split.hpp#L152-L228)

```c
static unique_ptr create(table_view const& source_table,
size_t user_buffer_size, cudaStream_t stream, rmm::mr::device_memory_resource* mr = nullptr) //0
bool has_next() const //1
size_t next(device_span user_buffer) //2
vector build_metadata() const //3
size_t get_total_contiguous_size() const //4
```

## Limitations

* Inherently chunked packing/ contiguous splitting is a copy engine operation. But currently this is done by 1MB chunk-sized kernel operations using the streaming multiprocessors.
* All 1MB work loads are predetermined at the start of the operation. Due to this, during `chunked_pack` ctr (0), we need to provide a buffer size >= 1MB. I think it will be inefficient to provide a user buffer size that is not a multiple of 1MB.
* We need to provide a `user_buffer` >= provided user buffer size for the (2) `next` method.

Eg: 4 column table with each column 2.5MB sized. Pick 2MB as the user buffer size. Then copies will be done as follows.

| \# | What runs together | \~Bytes copied |
| :---- | :---- | :---- |
| 1 | C0 batch0 \+ C0 batch1 | 2.0 MiB |
| 2 | C0 batch2 \+ C1 batch0 | 1.5 MiB |
| 3 | C1 batch1 \+ C1 batch2 | 1.5 MiB |
| 4 | C2 batch0 \+ C2 batch1 | 2.0 MiB |
| 5 | C2 batch2 \+ C3 batch0 | 1.5 MiB |
| 6 | C3 batch1 \+ C3 batch2 | 1.5 MiB |

Even if I pass a 4MB user buffer to the `next` method, it would still copy in 6 iterations, because iterations are predetermined based on the column boundaries and `user_buffer_size`.
Ideally, using the copy engine, we can do this based on the `device_span` size provided in each iteration, and fill it fully (especially using `cudamemcpybatch`.

## Streaming multiprocessor vs copyengine

Even though packing is predominantly a copy operation, there could be some edge cases that require kernel launches.

- If the column views in the table are sliced (sliced from a parent column), then validity buffers and offset buffers need to be rebased (either have to be bitshifted or offsets need to be rebased to start from 0\)

# Suggested new API

```c
chunked_pack(table_view const& source_table, cudaStream_t stream) //0
bool has_next() const //1
size_t next(device_span user_buffer) //2
vector build_metadata() const //3
size_t get_total_contiguous_size() const //4
```

- No `user buffer size` limitation.
- Any sized span can be passed to (2). It will try to pack as much as possible. So, it returns how many bytes were written out from the beginning (ie. offset to the output packed buffer from that copy operation).

## Implementation

### Phase 1 - Unsliced columns

- Traverse all buffers in `source_table` and collect data ptrs and sizes during (0).
- Track current buffer idx, and buffer offset.
- For each `next` call, copy data to `user_buffer` until it's exhausted. Advance the current buffer idx and buffer offset accordingly. Return how many bytes written to the output buffer (with respect to the final packed buffer)
- Each of those copies can be done from a single `cudamemcpybatchasync` call

### Phase 2 - Support sliced columns

- In (2), if there are any sliced columns then,
- Validity bitmaps \- If the bitmap starts from the middle of the byte, copies will need to be bitshifted after the copy (May be done from thrust utils?)
- Offset buffer \- Offset needs to be rebased by subtracting the start offset. (Thrust again?)

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context, code examples, or references to existing implementations about the feature request here.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.