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

README.md

@@ -108,17 +108,24 @@ Implemented:
   kernel-specific panel formats and contiguous fast paths.
 - Microkernels: `ScalarKernel` (reference), `SIMDKernel` (default for real
   eltypes), `PlanarKernel` (split-complex, default for complex eltypes),
-  `OneMKernel` (the 1m method) and `FMAddSubKernel` (interleaved complex, used
-  for small-M complex contractions on AVX-512). Any of them can be named via
+  `OneMKernel` (the 1m method), `FMAddSubKernel` (interleaved complex, used
+  for small-M complex contractions on AVX-512) and `ComplexRealKernel`/
+  `RealComplexKernel` (a complex times a real operand on the real kernel,
+  default for those). Any of them can be named via
   `plan_contract(...; kernel = ...)`.
 - Register shapes derived from the detected ISA, demoted when `M` is too short
   for a full tile or when a smaller tile keeps stores into `C` unit-stride.
 - A BLIS five-loop (`NC`/`KC`/`MC`) nest reusing packed panels, with blocking
   from an analytical model of the detected cache sizes (fixed constants when
   they are undetected); plus automatically selected paths for dot-product,
   outer-product and read-B-in-place shapes.
-- `Float32`/`Float64`/`ComplexF32`/`ComplexF64`; `conjA`/`conjB` and each
-  operand's `StridedView.op` are applied during packing.
+- `Float32`/`Float64`/`ComplexF32`/`ComplexF64`, mixed freely (a complex input
+  needs a complex output). The compute type is the promoted eltype, or the
+  precision chosen with `accumulator = Float32`/`Float64`
+  (`plan_contract(...; accumulator)`, `QuasiStridedBackend(; accumulator)`);
+  operands are converted during packing and the result is rounded to
+  `eltype(C)` once. `conjA`/`conjB` and each operand's `StridedView.op` are
+  applied during packing.
 - Zero steady-state allocation on Julia >= 1.11 for `execute!` on a reused
   plan and for `tensorcontract!` through the backend (which pools workspaces
   per task). On Julia 1.10 `SIMDKernel`'s accumulator is not kept in
@@ -127,11 +134,10 @@ Implemented:
 Not implemented:
 
 - Anything but contraction: `tensoradd!`/`tensortrace!` under this backend are
-  forwarded to `StridedNative()`, so mixed `@tensor` networks still run.
-- Mixed eltypes (including real with complex; promotion is TensorOperations'
-  `promote_contract` job), other eltypes, non-strided operands, an output
-  aliased with an input, and a conjugated output view: `tensorcontract!` throws
-  an `ArgumentError` instead of falling back to another backend.
+  forwarded to `StridedNative()`, so whole `@tensor` networks still run.
+- Other eltypes, a complex input with a real output, non-strided operands, an
+  output aliased with an input, and a conjugated output view: `tensorcontract!`
+  throws an `ArgumentError` instead of falling back to another backend.
 - Being the default backend: `StridedBLAS()` is faster on most shapes.
 - Batch axes, traces/diagonals, threading, GPU execution, the 3m complex
   method, autotuning.

benchmark/bench_mixed.jl

@@ -0,0 +1,62 @@
+# Mixed-domain GEMM, single core: complex x real (CR) and real x complex (RC)
+# through the mixed-domain kernel, the planar kernel (the real operand promoted)
+# and the real GEMM of equal FMA count (2M x N x K for CR, M x 2N x K for RC).
+# The inner real tile is the host's default real shape, so `real/mixed` is the
+# fraction of real efficiency reached.
+#
+#   julia --project=benchmark benchmark/bench_mixed.jl [--dtypes ComplexF64,ComplexF32]
+#       [--sizes 512,2048] [--reps 5] [--outdir DIR]
+
+include(joinpath(@__DIR__, "harness.jl"))
+
+using QuasiStrided: ComplexRealKernel, RealComplexKernel, PlanarMethod, target_profile,
+    _derived_shape, _kernel_from_shape
+
+const RUN_DTYPES = parse_dtypes(argopt("dtypes", "ComplexF64,ComplexF32"))
+const SIZES = parse_ints(argopt("sizes", "512,2048"))
+const REPS = argopt("reps", 5)
+const OUTDIR = outdir()
+const CSV_PATH = joinpath(OUTDIR, "bench_mixed.csv")
+
+function time_plan(C, A, B; kw...)
+    plan = plan_contract(StridedView(C), StridedView(A), (1, 2), StridedView(B), (2, 3), (1, 3); kw...)
+    return median_time_s(() -> execute!(plan, 1, 0); reps = REPS)
+end
+
+function run(io)
+    print_env_header(stdout, "bench_mixed.jl")
+    println(io, "dtype,case,n,kernel,t_mixed,t_promoted,t_real")
+    rng = MersenneTwister(1)
+    for T in RUN_DTYPES, n in SIZES
+        R = real(T)
+        MRr, NRr, W = _derived_shape(target_profile(), R)
+        planar = _kernel_from_shape(_derived_shape(target_profile(), T), T, PlanarMethod())
+        cases = (
+            (
+                "CR", ComplexRealKernel(Val(MRr ÷ 2), Val(NRr), T, Val(W)),
+                randn(rng, T, n, n), randn(rng, R, n, n), (2n, n),
+            ),
+            (
+                "RC", RealComplexKernel(Val(MRr), Val(NRr ÷ 2), T, Val(W)),
+                randn(rng, R, n, n), randn(rng, T, n, n), (n, 2n),
+            ),
+        )
+        for (case, kernel, A, B, (Mr, Nr)) in cases
+            C = zeros(T, n, n)
+            tm = time_plan(C, A, B; kernel)
+            tp = time_plan(C, A, B; kernel = planar)
+            tr = time_plan(zeros(R, Mr, Nr), randn(rng, R, Mr, n), randn(rng, R, n, Nr))
+            tag = "$(mr(kernel))x$(nr(kernel))/W$(lanewidth(kernel))"
+            @printf(
+                "%-10s %s %5d  %-10s mixed %8.4f s  promoted %8.4f s  real %8.4f s  promoted/mixed %.2f  real/mixed %.2f\n",
+                T, case, n, tag, tm, tp, tr, tp / tm, tr / tm
+            )
+            println(io, join((T, case, n, tag, tm, tp, tr), ','))
+            flush(stdout)
+        end
+    end
+    return
+end
+
+open(run, CSV_PATH, "w")
+println("wrote ", CSV_PATH)

src/QuasiStrided.jl

@@ -33,6 +33,7 @@ include("microkernels/simd.jl")
 include("microkernels/planar.jl")
 include("microkernels/onem.jl")
 include("microkernels/fmaddsub.jl")
+include("microkernels/mixed.jl")
 
 # --- Planning: labels, conjugation, kernel and blocking choice, the plan ---
 include("planning/labels.jl")
@@ -72,6 +73,7 @@ end
             :public, :contract!, :plan_contract, :execute!, :ContractPlan,
             :ContractWorkspace, :Blocking, :default_blocking,
             :ScalarKernel, :SIMDKernel, :PlanarKernel, :OneMKernel, :FMAddSubKernel,
+            :ComplexRealKernel, :RealComplexKernel,
             :target_profile, :cache_topology,
             :TargetProfile, :CacheLevel
         )

src/execution/barrier.jl

@@ -48,28 +48,32 @@ struct _NestPath{UNPACKED_B, AFF, SPLIT} end
 
 # N's B flag stays `false` when B is read in place: that path builds no B
 # sliver axis, so the flag would only split specialisations. A split group is
-# enumerated in another order, so its ramp flags are `false`.
+# enumerated in another order, so its ramp flags are `false`. `panel`: C is
+# `_PanelPath`'s dense panel, whose M and N maps are ramps.
 @inline function _nest_path(
         unpacked_b::Bool, mgroup::AxisGroup, ngroup::AxisGroup, kgroup::AxisGroup,
-        split_a::Bool, split_b::Bool
+        split_a::Bool, split_b::Bool, panel::Bool = false
     )
     split_b &= !unpacked_b
     return _nest_path_from_flags(
-        unpacked_b,
-        !split_a && _is_ramp_map(mgroup, 1), !split_a && _is_ramp_map(mgroup, 2),
-        !unpacked_b && !split_b && _is_ramp_map(ngroup, 1), !split_b && _is_ramp_map(ngroup, 2),
+        panel ? _PANEL_PATHS : _NEST_PATHS, unpacked_b,
+        !split_a && _is_ramp_map(mgroup, 1), !split_a && (panel || _is_ramp_map(mgroup, 2)),
+        !unpacked_b && !split_b && _is_ramp_map(ngroup, 1), !split_b && (panel || _is_ramp_map(ngroup, 2)),
         _is_ramp_map(kgroup, 1), _is_ramp_map(kgroup, 2), split_a, split_b
     )
 end
 
+# The nest run against a dense panel of C in the compute type.
+struct _PanelPath{P <: _NestPath} end
+
 # Table lookup rather than `_NestPath{u, aff, split}()` built at runtime (a
 # dynamic type application) or a 512-leaf branch tree (slow to infer).
-@inline function _nest_path_from_flags(flags::Vararg{Bool, 9})
+@inline function _nest_path_from_flags(table::Vector{Any}, flags::Vararg{Bool, 9})
     index = 1
     for i in 1:9
         index += flags[i] << (i - 1)
     end
-    return @inbounds _NEST_PATHS[index]
+    return @inbounds table[index]
 end
 
 const _NEST_PATHS = let
@@ -81,6 +85,8 @@ const _NEST_PATHS = let
     table
 end
 
+const _PANEL_PATHS = Any[_PanelPath{typeof(p)}() for p in _NEST_PATHS]
+
 # The dot path at lane width `W`; `MATB`: the matrix operand is B (`Qm == 1`).
 struct _DotPath{MATB, W} end
 

src/execution/dot.jl

@@ -174,7 +174,7 @@ end
         gptr::Ptr{T}, vstorage::SV, vbase::Int, vkbuf::Vector{Int}, kblock::Int, transform::F
     ) where {T, SV, F}
     @inbounds for t in 0:(kblock - 1)
-        z = transform(vstorage[vbase + vkbuf[t + 1] + 1])::T
+        z = convert(T, transform(vstorage[vbase + vkbuf[t + 1] + 1]))
         unsafe_store!(gptr + sizeof(T) * t, z)
         if T <: Complex
             unsafe_store!(gptr + sizeof(T) * (kblock + t), Complex(imag(z), real(z)))