WIP: batched out-of-place tensor permutation

[lkdvos]

WIP — not ready for review, needs cleanup before merging.

Adds batched_permutedims!/plan_batched_permutedims: permute a batch of tensors sharing one perm in a single CPU pass or GPU kernel launch, with selectable GPU strategies (elementwise, thread-tile, cooperative shared-memory tile) and optional per-tensor alpha/beta scaled accumulation.

Still to do before this is mergeable: general read-through for remaining rough edges, decide on BP_AUTO's default strategy per family, and reconcile with upstream main.

Benchmarks (RTX A6000, batched transposes; achievable ceiling via plain copyto! ≈ 640/660 GB/s F32/F64, not the ~768 GB/s spec figure)

GB/s (min-of-20, plan_batched_permutedims built once and reused):

batch	strategy	F32	F64
200 tensors, (48,48,24)	BP_ELEMENTWISE (default)	113	222
	BP_THREADTILE	195–375 (orientation-dependent)	250–437
	BP_GROUPTILE	409	515
1 large (2048²) + 400 tiny (8×8)	BP_ELEMENTWISE (default)	110	215
	BP_THREADTILE	113–130	120–172
	BP_GROUPTILE	195	306
2000 tiny (16×16×16)	BP_ELEMENTWISE (default)	58	105
	BP_THREADTILE	74–79	110–129
	BP_GROUPTILE	90	152
CPU (same 3 batches, in order)	—	22 / 7 / 27	17 / 3 / 25

BP_GROUPTILE wins on every one of these transpose batches (1.6–4x the default), but only applies to the transpose family; BP_THREADTILE is the better fit for non-transpose (payload) permutations. BP_AUTO conservatively defaults to BP_ELEMENTWISE for now — picking a smarter per-family default is one of the open items above. The "many tiny tensors" case (last row) is dominated by per-call host-side bookkeeping rather than kernel time; that overhead was a specific target of later work in this branch's history and is no longer the main bottleneck it once was.

Scaled accumulation (alpha/beta): omitting them costs nothing (identical to the plain copy). With beta == 0, overhead is ~1–4%. With beta != 0 (destination read+accumulate) on a transpose, BP_ELEMENTWISE costs ~3.8–4x the unscaled call, while BP_GROUPTILE costs only ~1.5x and is faster in absolute terms.

🤖 Generated with Claude Code

[kshyatt]

Could we get some AMD benchmarks as well?

[lkdvos]

I don't have AMD access unfortunately

[kshyatt]

I do if you can send me the script you used for the benchmarks above!