Porting src/RSE/UoE/jit_comparison/src/jit_comparison/intro_numba.py to Julia.
Numba needed @jit to opt a function into compilation to machine code. Julia has no such decorator: every function is compiled to machine code, specialised on its argument types, the first time it is called with a new combination of types — the whole language works the way mul_numba did. So instead of contrasting an interpreted function with a compiled one, three compiled functions below compute the same thing and differ only in how much of the fusion work they leave to the compiler versus doing it themselves by calling straight into BLAS. All three share LinearAlgebra.mul!’s five-argument signature, (C, A, B, α, β), and all three write their result into C — that’s what makes the comparison fair: the naive and fused versions don’t get to allocate a fresh output array while mul_native! doesn’t.
mul_native! (generic function with 3 methods)mul_naive! reads like the eager NumPy version: * and + each build and return a whole new array — α * (A * B), β * C, and their sum are three separate temporaries — and only the final .= writes that last temporary into C. mul_fused! differs only by adding dots — Julia’s broadcast fusion syntax — all the way through, including the assignment. A * B stays plain * because matrix multiply isn’t elementwise (it lowers to a BLAS call, same as Numba’s A @ B); everything elementwise around it, including the final store into C, becomes dotted, and the dots are what tells the compiler to fuse. C appears on both sides of mul_fused!’s assignment — safe here because broadcasting into a destination that also appears in the broadcasted expression is computed and stored one element at a time, so there’s no danger of reading a location after it’s been overwritten.
mul_native! is C, A, B, α, β handed straight to LinearAlgebra.mul! — which computes exactly α·(A·B) + β·C, in place, because that five-argument signature is BLAS’s gemm calling convention. Nothing is fused here because there’s nothing left to fuse: the whole expression is one call into a library routine that never materialises A * B at all.
3.141592653589793using BenchmarkTools, Printf
mul_naive!(copy(C), A, B, α, β) # compile all three, once, before benchmarking
mul_fused!(copy(C), A, B, α, β)
mul_native!(copy(C), A, B, α, β)
trials = [
"mul_naive!" => (@benchmark mul_naive!($C, $A, $B, $α, $β) setup = (fill!($C, 0))),
"mul_fused!" => (@benchmark mul_fused!($C, $A, $B, $α, $β) setup = (fill!($C, 0))),
"mul_native!" => (@benchmark mul_native!($C, $A, $B, $α, $β) setup = (fill!($C, 0)))
]
@printf("%-12s %12s %12s %8s\n", "impl", "time", "memory", "allocs")
for (name, trial) in trials
t = median(trial)
@printf(
"%-12s %12s %12s %8d\n",
name,
BenchmarkTools.prettytime(t.time),
BenchmarkTools.prettymemory(t.memory),
t.allocs
)
endimpl time memory allocs
mul_naive! 633.376 ns 16.34 KiB 12
mul_fused! 277.267 ns 4.09 KiB 3
mul_native! 103.032 ns 0 bytes 0All three now mutate C, so all three need setup = ... — but resetting with fill!($C, 0) rather than Cc = copy($C) is deliberate: a fresh copy is itself an allocation, and while @benchmark’s setup expression runs outside the timed region, it would still leave every sample reading from last sample’s output instead of a known state, since it’s the same C object being mutated and re-read for β·C on every call. fill! zeros it in place, in O(1) extra memory, before each sample runs. The shape of the table is the point: memory roughly halves from mul_naive! to mul_fused!, then drops to zero for mul_native!, and time drops with it. Fusing the elementwise part removes one of two temporaries; skipping Julia’s arithmetic entirely and letting BLAS accumulate β·C in place removes the other.
Numba discovers the fusion by analysing typed IR after the fact. Julia decides it before any type is known at all — during lowering, the step that turns parsed syntax into the untyped IR the compiler works from. @code_lowered shows that IR directly.
mul_naive!, lowered:CodeInfo( 1 ─ %1 = Main.:+ │ %2 = Main.:* │ %3 = Main.:* │ %4 = dynamic (%3)(A, B) │ %5 = dynamic (%2)(α, %4) │ %6 = Main.:* │ %7 = dynamic (%6)(β, C) │ %8 = dynamic (%1)(%5, %7) │ %9 = dynamic Base.broadcasted(Base.identity, %8) │ dynamic Base.materialize!(C, %9) │ %11 = C └── return %11 )
mul_fused!, lowered:CodeInfo( 1 ─ %1 = Main.:+ │ %2 = Main.:* │ %3 = Main.:* │ %4 = dynamic (%3)(A, B) │ %5 = dynamic Base.broadcasted(%2, α, %4) │ %6 = Main.:* │ %7 = dynamic Base.broadcasted(%6, β, C) │ %8 = dynamic Base.broadcasted(%1, %5, %7) │ dynamic Base.materialize!(C, %8) │ %10 = C └── return %10 )
mul_naive!’s lowered IR is a straight-line chain of eager calls: *(A, B), then *(α, ...), then *(β, C), then +, then a final materialize! that stores the result into C — each one a call that runs to completion and hands back a full array before the next starts. mul_fused!’s IR calls Base.broadcasted instead of */+ for every dotted operator — and broadcasted doesn’t compute anything, it just records “multiply these, lazily.” The three broadcasted calls nest into a single unevaluated tree, and only the final Base.materialize! actually walks it, in one pass, writing the result directly into C. The macro-expanded syntax already encodes “one loop” before type inference, LLVM, or anything else in the pipeline has run — this is Julia’s dot syntax doing at parse time what Numba’s optimiser had to reconstruct from typed IR.
From here the pipeline is the same shape as Numba’s: typed IR, then LLVM IR, then assembly.
io = IOBuffer()
show(io, (@code_typed mul_fused!(C, A, B, α, β)))
typed_ir = String(take!(io))
typed_lines = split(typed_ir, '\n')
println(length(typed_lines), " lines of typed IR. The interesting bits:\n")
for l in typed_lines
if occursin("generic_matmatmul_wrapper", l) || occursin("mul_float", l) || occursin("add_float", l)
println(l)
end
end618 lines of typed IR. The interesting bits:
7 ─── %24 = invoke LinearAlgebra.generic_matmatmul_wrapper!(%20::Matrix{Float64}, 'N'::Char, 'N'::Char, A::Matrix{Float64}, B::Matrix{Float64}, true::Bool, false::Bool, $(QuoteNode(Val{LinearAlgebra.BlasFlag.GEMM}()))::Val{LinearAlgebra.BlasFlag.GEMM})::Matrix{Float64}
153 ─ %466 = intrinsic Base.mul_float(α, %459)::Float64
174 ─ %534 = intrinsic Base.mul_float(β, %527)::Float64
177 ─ %538 = intrinsic Base.add_float(%466, %534)::Float64A * B still calls out to LinearAlgebra.generic_matmatmul_wrapper! — the BLAS call didn’t fuse, exactly like A @ B in the Numba example. The mul_float/add_float pair is the fused α·x + β·y from the broadcast — one multiply-add per element, downstream of that call, in place of three separate array-producing operators.
io = IOBuffer()
code_llvm(io, mul_fused!, Tuple{typeof(C),typeof(A),typeof(B),typeof(α),typeof(β)})
llvm_ir = String(take!(io))
println(
count('\n', llvm_ir),
" lines of LLVM IR. A sample of the fused multiply-add sequence buried inside it:\n",
)
for line in Iterators.take(
(l for l in eachsplit(llvm_ir, '\n') if occursin("fmul", l) || occursin("fadd", l)),
10,
)
println(line)
end11554 lines of LLVM IR. A sample of the fused multiply-add sequence buried inside it:
%32 = fmul double %31, %"\CE\B1::Float64"
%34 = fmul double %33, %"\CE\B2::Float64"
%35 = fadd double %32, %34
%38 = fmul double %37, %"\CE\B1::Float64"
%40 = fmul double %39, %"\CE\B2::Float64"
%41 = fadd double %38, %40
%44 = fmul double %43, %"\CE\B1::Float64"
%46 = fmul double %45, %"\CE\B2::Float64"
%47 = fadd double %44, %46
%50 = fmul double %49, %"\CE\B1::Float64"<8 x double> — LLVM auto-vectorised the fused loop, same as in the Numba example, just with a wider register (AVX-512 on this machine: 8 doubles at a time, versus 4 via AVX2). Most of the thousands of surrounding lines are bounds checks and array-header bookkeeping, not arithmetic.
io = IOBuffer()
code_native(io, mul_fused!, Tuple{typeof(C),typeof(A),typeof(B),typeof(α),typeof(β)})
asm = String(take!(io))
println(
count('\n', asm),
" lines of x86-64 assembly. A sample of the vectorised multiply/add instructions:\n",
)
vector_ops =
(l for l in eachsplit(asm, '\n') if occursin("mulpd", l) || occursin("addpd", l))
for line in Iterators.take(vector_ops, 10)
println(line)
end8633 lines of x86-64 assembly. A sample of the vectorised multiply/add instructions:
vmulpd zmm3, zmm0, zmmword ptr [r14 + 8*rbx - 128]
vmulpd zmm9, zmm1, zmmword ptr [r15 + 8*rbx - 128]
vmulpd zmm2, zmm0, zmmword ptr [r14 + 8*rbx - 192]
vmulpd zmm6, zmm1, zmmword ptr [r15 + 8*rbx - 192]
vmulpd zmm4, zmm0, zmmword ptr [r14 + 8*rbx - 64]
vmulpd zmm10, zmm1, zmmword ptr [r15 + 8*rbx - 64]
vmulpd zmm5, zmm0, zmmword ptr [r14 + 8*rbx]
vaddpd zmm3, zmm3, zmm9
vmulpd zmm9, zmm1, zmmword ptr [r15 + 8*rbx]
vaddpd zmm2, zmm2, zmm6vmulpd/vaddpd on zmm registers: the same fused arithmetic, now as vector instructions.
mul_native! collapses the towermul_fused! still has a tower to walk down — lowered IR, typed IR, LLVM, assembly — because Julia has to generate the fused loop. mul_native! doesn’t:
CodeInfo( 1 ─ %1 = invoke LinearAlgebra.generic_matmatmul_wrapper!(C::Matrix{Float64}, 'N'::Char, 'N'::Char, A::Matrix{Float64}, B::Matrix{Float64}, α::Float64, β::Float64, $(QuoteNode(Val{LinearAlgebra.BlasFlag.GEMM}()))::Val{LinearAlgebra.BlasFlag.GEMM})::Matrix{Float64} └── return %1 ) => Matrix{Float64}
One line: α and β pass straight through to generic_matmatmul_wrapper!, which forwards them to BLAS’s gemm — the same routine A * B already called, just given the accumulate coefficients it was always able to take. There’s no elementwise Julia code left to lower, type-infer, or vectorise, so LLVM IR and assembly for this function are uninteresting: it’s a calling-convention wrapper around a ccall, nothing more.
Three rungs of the same ladder, all writing into the same C: mul_naive! leaves three temporaries for the allocator before its final store, mul_fused! collapses the elementwise half into one loop and gets down to zero temporaries of its own, mul_native! hands the whole expression to a library that never needed a temporary in the first place. It’s the same trick — fuse operators to avoid materialising intermediates — applied at three different levels: by hand, by the language’s broadcast syntax, and by the numerical library underneath both. XLA does the same fusion at a much larger scale, later in this talk.