Skip to content

Broaden #1003 into a shared CPU/GPU elementwise framework (follow-up to #1067) #1069

Description

@yingjerkao

Context

Follow-up to PR #1067, which converts GPU elementwise arithmetic to typed
std::visit dispatch and removes the legacy per-dtype tables. In his review of
#1067, @ianmccul noted that the non-contiguous mapper kernels introduced/carried
forward there should be treated as transitional code, and that the broader
work tracked in #1003 should be broadened to a shared CPU/GPU elementwise
framework. This issue collects the concrete items from that review so they aren't
lost.

The typed-dispatch cleanup and dtype-correctness fixes in #1067 are independently
worthwhile and can land first; the items below are the follow-up.

Architecture (broaden #1003)

Move to a single shared CPU/GPU elementwise framework:

  • common operation functors and dtype/promotion policies (the functor is already
    shared within one GPU arithmetic family, but traversal is not),
  • common host-side stride/layout preprocessing into one compact,
    trivially-copyable layout descriptor,
  • backend-specific CPU/CUDA execution that consumes the same descriptor.

This removes the traversal duplication that currently exists across
cuArithmeticDispatch.cuh, cuiArithmeticDispatch.cuh, cuCpr_dispatch.cu, and
the analogous CPU traversal from #1056.

Concrete items

  1. Rank-dependent shared-memory launch limit (correctness). The non-contiguous
    kernels request 512 * rank * sizeof(uint64_t) = 4096 * rank bytes of dynamic
    shared memory for per-thread coordinate vectors. Rank 13 already exceeds the
    common 48 KiB limit, so high-rank non-contiguous arithmetic fails at launch with
    no rank check or fallback. (Pre-existing: inherited from the legacy kernels.)
    Fix by not putting the per-thread coordinate vector in dynamic shared memory, or
    at minimum add an explicit supported-rank check / rank-derived block size.

  2. Pass layout metadata as a kernel argument, not via managed allocations. Each
    non-contiguous op currently does 5 cudaMallocManaged allocations + 2 copies +
    launch + 5 frees for a few hundred bytes of metadata. Pack extents/strides/
    mappers into one trivially-copyable struct passed by value (e.g.
    __grid_constant__ LayoutMetadata with std::array<uint64_t, MAX_RANK> fields;
    ~1.3 KiB at rank 32, well under the 32 KiB parameter limit on modern devices),
    or copy once with cudaMemcpyAsync / prefetch. (Gemini's "host write to device
    pointer" flag was a false positive — cuCalloc_gpu is cudaMallocManaged — but
    demand-paging tiny managed allocations is still a poor metadata path.)

  3. In-place: iterate LHS physical order; only map the RHS. For in-place ops the
    LHS is also the output, so traversing LHS physical order guarantees coalesced
    LHS reads and output writes; only the RHS may need an index transform. Compose
    the two permutations host-side into a single LHS-physical -> RHS-physical map.
    The current kernel redundantly decodes a logical index, computes Lidx, then
    Ridx, and writes out[Lidx].

  4. Same-layout fast path. When Lt.shape() == Rt.shape() && Lt.invmapper() == Rt.invmapper(), corresponding logical elements share physical offsets, so the op
    is a plain physical-buffer zip (out[i] = op(lhs[i], rhs[i])) with no mapper --
    true even for non-commutative ops. Today the front end takes the full mapper path
    whenever either operand is individually non-contiguous, so two identically
    permuted operands pay the whole mapper cost unnecessarily.

  5. Uniform non-contiguous reachability. Out-of-place GPU Mul/Div/Mod/Cpr
    front ends still reject (or contiguous-ize) non-contiguous tensor/tensor
    operands; only Add/Sub and the in-place paths reach the mapper-aware kernels.
    The shared framework should make the mapper-aware kernel uniformly reachable.

  6. Benchmarks to guide the redesign. Before treating the mapper path as final,
    benchmark (over multiple sizes/ranks): contiguous a += b; identically permuted
    a += b (should match contiguous once the fast path exists); contiguous LHS +
    permuted RHS; differently permuted operands; and contiguous()-copy-then-zip
    including copy cost.

Credit: findings and suggestions by @ianmccul (posted via OpenAI Codex) on #1067.
Related: #1003, #1067.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions