Skip to content

[FEA] Optimize variant extraction kernels by packing optional<size_type>/op_status return values into a single 64-bit int #23656

Description

@abigalekim

Is your feature request related to a problem? Please describe.

Several __device__ helpers in cpp/src/io/parquet/experimental/variant_extract.cu (e.g. find_key_in_metadata, defined in variant_extract.cu on line 234) return cuda::std::pair<cuda::std::optional<size_type>, op_status>. cuda::std::optional<size_type> uses 8 bytes to hold a 4-byte payload plus a 1-bit validity flag, and pairing it with a 1-byte op_status adds further padding. In a register-pressure-sensitive device kernel that runs per-row over potentially large batches, there is a potential that this representation of data could result in worse occupancy.

Describe the solution you'd like

Investigate whether encoding (value, valid, op_status) into a single packed 64-bit integer would reduce register usage and improve occupancy for these kernels.

So encode would be:

val | (op_status << 33) | (valid << 32)

And decode would be:

val = packed & 0xFFFFFFFF;
valid = (packed >> 32) & 0x1;
op_status = (packed >> 33) & 0xFF;

We would apply this consistently across all the different functions that have this potential issue (find_key_in_metadata, locate_object_field, locate_array_element, resolve_path, etc.)

Describe alternatives you've considered

We could leave the current optimal/pair-based encoding as-is, depending on whether these functions show a register occupancy bottleneck when profiled.

Additional context

Raised during review of #23560. The comment where this is addressed is #23560 (comment).

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions