Skip to content

Mixed element types and mixed-domain (real × complex) kernels - #12

Merged
lkdvos merged 29 commits into
mainfrom
mixed
Sep 30, 2026
Merged

lkdvos merged 29 commits into
mainfrom
mixed

Conversation

@lkdvos

@lkdvos lkdvos commented Sep 29, 2026

Copy link
Copy Markdown
Member

Mixed element types and mixed-domain kernels

plan_contract, contract! and QuasiStridedBackend now accept any combination of Float32, Float64, ComplexF32 and ComplexF64 operands. The one exception is a complex input into a real C.

API

  • accumulator = nothing | Float32 | Float64 keyword on plan_contract and contract!.
  • QuasiStridedBackend(; accumulator = nothing). The backend type is now QuasiStridedBackend{A}.
  • New public kernels ComplexRealKernel and RealComplexKernel, which can be named via kernel = ....
  • ContractWorkspace{T,VT,PT} / ContractPlan{…,VT,PT}: the new parameter PT is the storage type of the compute-type panel of C.

Semantics

  • Compute type:
    • Default: T = promote_type(TA, TB, TC).
    • With accumulator: that precision, in the promoted domain.
  • Operands are converted to T while packing; the contiguous pack fast paths convert whole lanes.
  • alpha*AB + beta*C is evaluated in T and rounded to eltype(C) once, as in cuBLAS GemmEx.
  • When eltype(C) is narrower than T and K spans more than one kc block, a dense M × min(N, nc) panel of eltype T in the workspace carries the partial sums across kc blocks. This keeps "rounded once" true. The panel is owned by the allocator and freed by release!.

Design

  • ComplexRealKernel runs SIMDKernel{2MR,NR,real(T),W} on interleaved A × real B. It uses the fmaddsub accumulator layout, reuses the fmaddsub stores, and is eligible for unpacked B.
  • RealComplexKernel runs SIMDKernel{MR,2NR,real(T),W} on real A × interleaved B. It uses the planar stores through a column-pair index map.
  • Both cost 2 real FMAs per complex multiply-add, instead of 4 when the real operand is promoted.
  • The kernel is selected statically from the operand eltypes. Shapes are the real default for real(T) on the equivalent real problem, mapped back, so the extent, store, small-M and run demotions carry over.
  • Only the chosen kernel is compiled. Same-eltype code_llvm is unchanged relative to the pre-feature base for execute!, the nest, stores, packs and planning; the commit messages list what was compared.

Measurements

Slurm, bench_mixed.jl, single thread, at 7f68104's parent 1558b71 (the kernel code is unchanged since):

node promoted / mixed real GEMM / mixed
Ice Lake (Xeon 8362) 1.94–2.13× 0.98–1.01
Genoa (EPYC 9474F) 1.94–2.07× 0.98–1.02
  • The table covers C64 and C32, CR and RC, at 512³ and 2048³. Example: C64 CR at 2048³ on Ice Lake takes 0.360 s mixed, 0.699 s promoted and 0.354 s for the real GEMM.
  • Workstation measurements (indicative only):
    • Converting packs now run at 1.0–1.4× the time of a same-type pack, down from 2.2–4.9× on the generic path.
    • Panel precision: a Float32 C with accumulator = Float64 is correctly rounded across kc blocks; per-block rounding was off by up to about 4 ulp.
    • Panel cost: 1024³ with the panel takes about 1.04× the time of an all-Float64 run.
    • First-call latency (probe_ttfx.jl) is unchanged: load 0.5 s, first Float64 call 3.0–3.1 s, first ComplexF64 call 2.0–2.1 s. A first mixed @tensor call takes 1.04–1.11× as long as the ComplexF64 one.
    • Full test suite: about 830 s, within noise of the baseline.

Tests

  • test/execution/test_mixed.jl:
    • Compute type, kernel and path selection over 9 rows varying eltype, layout, conj and path.
    • Rejection of a named kernel whose domain doesn't match the operands.
    • A deterministic round-once table: several jc blocks under ManualAllocator, packed B, a complex panel with RealComplexKernel and strided C, the TensorOperations route, M == 1, and beta = 0 with NaN in C.
    • Zero allocations in the steady state.
  • test/microkernels/test_mixed_kernel.jl: the kernel contract suite for both mixed-domain kernels, packing from storage of another precision, and end to end.
  • Lane-conversion rows in the existing pack and store fast-path tests, plus selection tests, rejection tests and mixed @tensor cases.

Known limitations

  • No complex input into a real C. Only Float32 and Float64 accumulators.
  • The tile-wise oracle execute_tilewise! still rounds a narrow C once per kc block.
  • The panel needs M * min(N, nc) elements of T in the workspace. The task-local workspace pool of the TensorOperations backend only grows, so it keeps the panel after the call.
  • On AVX2 and NEON the complex vector fast paths are gated off, so the mixed kernels use scalar stores there.
  • Not optimized yet: real × real into a complex C, a row-major C for the mixed kernels (no swap between the two kernels), and the dot/outer fast paths for mixed domains.

🤖 Generated with Claude Code

lkdvos and others added 29 commits September 28, 2026 18:33
The contiguous packers (real A, complex planar/interleaved/1e) and the real,
planar and fmaddsub vector stores now accept dense rank-1 storage of any
supported type in the kernel's domain (Float32/Float64 for real kernels,
ComplexF32/ComplexF64 for complex ones). Lanes load in the storage real type,
convert to the kernel real type, and C rounds back once on store. Eligibility
folds statically on the storage type.

Same-eltype specializations emit identical LLVM IR to 42f962b (59
specializations of store_tile!, _store_tile_{vector,planar_vector,
fmaddsub_vector}!, _pack_a_contiguous!, _pack_complex_contiguous!,
unsafe_pack_a!/unsafe_pack_b! for F64/F32/C64/C32, diffed after normalizing
specialization suffixes).

Local pack_a! into a PackedPanel, kc=256, best of 15 (workstation, indicative):
  F64 <- F32 (MR=16):   1370 -> 615 ns   (same-type F64: 633 ns)
  F32 <- F64 (MR=32):   3083 -> 903 ns   (same-type F32: 627 ns)
  C64 <- C32 planar:    1817 -> 639 ns   (same-type C64: 661 ns)
  C32 <- C64 planar:    3245 -> 909 ns   (same-type C32: 662 ns)
Full Pkg.test: 778 s wall. probe_ttfx (3 runs, after/base): load 0.5/0.5 s,
first call F64 3.0/2.9 s, C64 2.0/2.1 s, K=1 outer 0.5/0.5 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
plan_contract, contract! and QuasiStridedBackend accept any combination of
Float32/Float64/ComplexF32/ComplexF64 operands, except a complex input into a
real C. The compute type is promote_type(TA, TB, TC), or with
`accumulator = Float32/Float64` that precision in the promoted domain; the
plan, workspace and kernel are keyed on it. The scalar store helpers, the
unpacked-B loads, the dot gather and the outer-product B load convert to the
compute type explicitly; a real operand of a complex contraction is promoted
by the generic packer.

Same-eltype code is unchanged: code_llvm of _axpby_tile!/_axpby_at!/
scale_tile!/execute! (18 F64/C64 specializations) is identical to 8b16443
modulo name suffixes, and return types of plan_contract/_planned (automatic
and named kernel) and execute! are unchanged.

Partial sums still round-trip through C between kc blocks of K, so a C
narrower than the compute type rounds once per K block, not once overall.

Workstation under load (loadavg ~16), indicative only:
  Full Pkg.test, run concurrently with 8b16443: 1281 s vs 1250 s (+2.5%).
  probe_ttfx, 3 runs (new / 8b16443): load 0.5-0.6 / 0.5-0.6 s, first F64
  3.7-4.4 / 3.8-4.2 s, C64 2.3-3.0 / 2.7-3.0 s, K=1 outer 0.7 / 0.6-0.8 s.
  Fresh-process first @tensor, 3 runs: F64xF64 3.2-5.8 s, F32xF64->F64
  3.3-3.6 s, C64xC64 3.2-3.5 s, F64xC64->C64 3.1-3.6 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…IMD kernel

ComplexRealKernel runs SIMDKernel{2MR,NR,real(T),W} over interleaved A and
real B (accumulator in the fmaddsub layout, reusing its stores, unpacked B
eligible). RealComplexKernel runs SIMDKernel{MR,2NR,real(T),W} over real A and
interleaved B, storing through the planar generators with a column-pair index
map. The descriptor admits (Interleaved, Real) and (Real, Interleaved) for a
complex T; a RealFormat operand packs real(T). Named kernels run through
plan_contract unchanged.

Codegen (code_llvm, debuginfo stripped, name suffixes normalized) against
096ab7f: identical for store_tile!/execute_tile! of SIMD/planar/fmaddsub/1m,
_store_tile_planar[_vector]!, _store_tile_fmaddsub_vector!,
_store_tile_lanepair!, and unsafe_pack_a!/unsafe_pack_b! (dense/scattered,
identity/conj, PackedPanel/Vector, same and mixed storage precision): 264
functions. The mixed kernels' _accumulate_step is identical to the equivalent
SIMDKernel's.

Pkg.test() run concurrently with 096ab7f on a Cascade Lake workstation:
856 s vs 888 s wall.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…l GEMM

Indicative single-core timings on a Cascade Lake workstation (noisy), inner
tile = default real shape (F64 32x6/W8, F32 64x6/W16):

  dtype       case     n   promoted/mixed  real/mixed
  ComplexF64  CR     512   1.84            0.93
  ComplexF64  RC     512   1.94            0.95
  ComplexF64  CR    2048   1.76            1.00
  ComplexF64  RC    2048   2.01            1.05
  ComplexF32  CR     512   1.92            1.00
  ComplexF32  RC     512   2.07            0.95
  ComplexF32  CR    2048   1.91            1.00
  ComplexF32  RC    2048   1.99            1.02

probe_ttfx.jl 3x against 096ab7f: Float64 first call 3.2-3.3 s vs 3.1-3.2 s,
ComplexF64 2.1-2.2 s vs 2.1-2.2 s (within noise).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
When real(eltype(C)) is narrower than real(T) and K spans more than one
kc block, the nest runs each jc column block against a dense
M x min(N, nc) panel of T in the workspace (new grow-only slot c_panel):
the block of C is loaded into the panel (skipped for beta == 0), all
pc/ic/jr/ir iterations store into it, and it is rounded back to C once.
The choice is a _PanelPath{P} singleton from _select_path, predicted
exactly by the planning hint (the default kc does not depend on the
complex method), so the plain nest is never compiled for such a plan.
The dot path is skipped in that case; the outer path has one K block.
execute_tilewise! keeps rounding per kc block.

Float32 operands, accumulator = Float64, Float32 C, 64x512x48 at the
default kc = 341: every element is now within 1 ulp of the correctly
rounded Float64 result (was up to ~4 ulp).

Same-eltype (F64, C64; packed, unpacked B) LLVM for execute!,
_execute_path! and _execute_nest! is identical to 096ab7f modulo
function-name suffixes and the workspace's dereferenceable size
(168 -> 176 bytes, the new field); return types unchanged.
Local measurements (noisy workstation, indicative): Pkg.test() 879.6 s
vs 873.2 s for 096ab7f run concurrently; probe_ttfx.jl x3 identical
(Float64 3.0-3.1 s, ComplexF64 2.0-2.1 s, outer 0.5 s); first
@tensor call with Float32 operands and accumulator = Float64 over two kc
blocks 3.35-3.42 s vs 3.17-3.28 s same-eltype Float64 (~1.05x).
1024^3, execute! only: narrow C with accumulator = Float64 27.9-28.6 ms
vs all-Float64 26.6-28.0 ms (~1.04x); 096ab7f's per-block-rounding
narrow C was 23.2-24.1 ms.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e real operand

A complex T with a complex A and a real B now plans a ComplexRealKernel,
a real A and a complex B a RealComplexKernel, statically from the operand
eltypes (no A/B swap for a complex T). Both run the real default shape
logic of real(T) on the real problem, (2Qm, Qn, 2run) for CR and
(Qm, 2Qn, run) for RC, and map the shape to (MR/2, NR, W) and
(MR, NR/2, W): the extent, store, small-M and run demotions carry over and
reuse the real cached defaults. A named CR/RC kernel whose RealFormat side
is complex throws ArgumentError at plan time.

Same-eltype selection and blocking for F32/F64/C32/C64 are unchanged over
7 shapes plus a swapped layout; code_llvm of _planned and the
plan_contract kwcall is identical to 293e212 modulo SSA/label names, type
tag addresses and GC-frame slot order (same frame size).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The default plan now selects the mixed-domain kernel, so the promoted
column names the host's planar kernel instead.

End-to-end, @tensor with QuasiStridedBackend(), median, workstation
(indicative only), this branch vs 293e212 (promoted):
  F64 x C64 -> C64  512^3:  6.2 ms vs 12.9 ms (2.08x); 2048^3: 0.415 s vs 0.804 s (1.94x)
  C64 x F64 -> C64  512^3:  6.9 ms vs 12.4 ms (1.80x); 2048^3: 0.456 s vs 0.798 s (1.75x)
First @tensor call, fresh process (3 runs): F64 x C64 3.17-3.30 s,
C64 x F64 3.40-3.54 s, C64 x C64 3.06-3.14 s (1.04x / 1.11x).
probe_ttfx.jl 3x: F64 3.0-3.1 s, C64 2.0-2.1 s, both trees.
Pkg.test() run concurrently with 293e212: 850.9 s vs 862.2 s wall.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The panel was a plain Vector{T} whatever the allocator. It now comes from
tensoralloc as a temporary like the packed panels (a PtrArray under
ManualAllocator, so no GC allocation) and release! frees it first.
ContractWorkspace and ContractPlan gain a type parameter for its storage.

Same-eltype code_llvm of execute!, _execute_path!, _execute_nest! and the
micro-tile specialisations (F64, F32, C64, C32, permuted C, dot, outer) is
identical to 565ad27 after normalising symbol numbering; code_typed
differs only in the printed workspace type.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
QuasiStridedBackend{Float16}() skipped the keyword constructor's check and
failed later with a MethodError. _compute_type now takes only Float32 and
Float64 accumulators and throws an ArgumentError otherwise. The backend
calls it first instead of repeating its eltype and complex-into-real
checks, so a rejected call still never touches the workspace pool.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The TO route predicts the path at the default kc. _hint_holds now also
requires that the plan's own blocking needs the panel for a hinted
_PanelPath (with a workspace panel large enough) and does not need it for
a hinted _NestPath; otherwise the plan's own path runs behind the barrier.
_c_panel_length is shared with plan sizing.

For a same-eltype plan the check folds away: code_llvm of _hint_holds and
_execute_hinted! (F64, F32, C64, C32, permuted C, dot, outer) is identical
to 565ad27 after normalising symbol numbering.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_lane_convertible(T, T) holds for every supported T, so the
DenseVector{T} test added nothing. Same-eltype code_llvm of the nest
(packing inlined) and the micro-tile specialisations for F64 and C64 is
unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Any other kernel type is now an ArgumentError at generator time instead
of silently getting the column-pair layout. The PlanarKernel branch is
unchanged, so the planar store is generated as before.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pack and end-to-end testsets now use the mk_contract shapes, so the
file builds 6 distinct mixed-domain kernels instead of 10. File time
(after the scalar/simd/planar kernel files, workstation) 24.7 s -> 24.5 s:
the six mk_contract shapes take about 15.6 s of it and stay.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Replaces the statistical comparison with the all-Float32 result. Along K
the contributions 1, 2^-30 and -1 fall in separate kc blocks, so the exact
2^-30 (plus a small beta*C) survives only if nothing is rounded to
eltype(C) between blocks. Six rows: several jc blocks under
ManualAllocator (checking the panel is allocator-owned), packed B, a
complex panel on the RC kernel with strided C, the TensorOperations hint
route at the default kc, M == 1, and beta = 0 on a NaN-filled C. With
_c_panel_needed forced to false every row fails.

The RC row of the main table now uses scattered operands.

File time (after test_execute.jl, workstation): 35.3 s -> 22.6 s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The descriptor already rejects a real T with an ArgumentError, and a
RealComplexKernel's inner SIMDKernel already checks that W divides MR.
The three-argument constructors forward to the four-argument ones.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e plan

A plan's workspace always holds the panel its own blocking needs
(_plan_contract sizes or grows it from the same inputs), so the hint only
has to agree with _c_panel_needed at the plan's kc. _c_panel_length is
inlined at its one remaining use, and the panel copy reads the workspace
and panel from the target plan.

Same-eltype code_llvm of execute!, _execute_path!, _execute_nest!,
_execute_hinted! and _hint_holds (F64, F32, C64, C32, permuted C, dot,
outer) is identical to 1558b71 after normalising symbol numbering.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_check_kernel_domain dispatches on ComplexRealKernel/RealComplexKernel
instead of forwarding through complex_method, and the run_m condition is
inlined.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Only the eltype check leaves it; the rename and the moved comment were
unrelated churn.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A real operand of a complex contraction is packed as real by
ComplexRealKernel/RealComplexKernel, not promoted, so "promoted at pack
time" and "Not implemented: dedicated real-times-complex kernels" were
wrong.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… mixed testset

The same-eltype defaults per ISA are covered by the real and complex
shape-selection testsets; the mixed testset keeps the mapping to the real
default of real(T), now also at a large extent, on the three synthetic
targets.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
test_mixed.jl checks steady-state allocations for the mixed-domain nest
and the panel path in one table. test_mixed_kernel.jl drops the checks of
a_reals/b_reals/accumulator_planes/complex_method (restating their
definitions; the selection tests check complex_method) and runs its
named-kernel end-to-end rows unconjugated and with both operands
conjugated; test_mixed.jl covers one-sided conjugation on the automatic
CR and RC plans.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…gnature

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
QuasiStridedBackend{Float16}() now throws the keyword constructor's
ArgumentError itself, so _compute_type no longer needs a throwing
fallback; plan_contract and contract! restrict the keyword's type.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…lidate

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Path marker: `_nest_path` takes both split flags and a `panel` flag; the
panel path is `_PanelPath{_NestPath{U, AFF, SPLIT}}` from a second
512-entry table, its C ramp flags `true` unless that side is split. Only
the chosen path is compiled.

Hint: `_hint_holds` still checks the panel decision; `_splits` and
`_split_path` now also cover a `_PanelPath` hint, so a split plan on the
TensorOperations route crosses the barrier to its own panel path.

Planning: `_pack_split` sees each operand's storage eltype size (lines,
page test, K walk) and whether its packed format is complex (the walk's
scatter cost), not `T`; for same-eltype plans these equal the previous
arguments. B is not split when C needs a panel: the panel holds `jc`
blocks in C's own N order, which a split N group does not enumerate
contiguously. A split under the panel is fine (M rows land at their
natural panel offsets).

Packing: `_pack_block_transposed!` converts to `_element_type(format, T)`,
so a real operand on a mixed-domain kernel's `RealFormat` side packs as
real (it was a MethodError in `panel_store!`). test_mixed.jl gains a
forced-split intensli_7 row on the RealComplexKernel.

Same-eltype code_llvm of execute!, _select_path, _execute_path!,
_execute_hinted!, _hint_holds and _execute_nest! (F64, C64: unpacked B,
packed B, forced A split) equals origin/main's after normalising
numbering, apart from the workspace's dereferenceable size.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
K = 700 spans more than one kc block only when the default kc is below 700,
which is false on targets with a larger cache model (the macOS aarch64 CI
runners), so the plan took the plain nest.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@lkdvos
lkdvos merged commit 4a77066 into main Sep 30, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant