By: Kristoffer Carlsson on Kristoffer Carlsson
Re-posted from: http://kristofferc.github.io/post/intrinsics/
Introduction
The good old days of processors getting a higher clock-speed every year have been over for quite some time now.
Instead, other features of CPUs are getting improved like the number of cores, the size of the cache, and the instruction set
they support.
In order to be responsible programmers, we should try our best to take advantage of the features the hardware
provides.
One way of doing this is multithreading which is exploiting the fact that modern CPUs have multiple cores and
can run multiple (possibly independent) instruction streams simultaneously.
Another feature to take advantage of is that a single core on modern processors can do operations on multiple values in one instruction (called SIMD – Single Instruction Multiple Data).
This article is intended to give a short summary of using SIMD in the Julia programming language.
It is intended for people already quite familiar with Julia.
The first part is likely familiar to people that have been using Julia for a while, the latter part, which is about explicitly calling
SIMD intrinsics might be new. Feel free to scroll down to the intrinsic bit or read the TLDR about it below.
TLDR (SIMD intrinsics)
To call an intrinsic like _mm_aesdec_si128:
- call the intrinsic from C
- use Clang with
-emit-llvmto figure out the LLVM intrinsic (for example) - call the intrinsic from Julia with
ccall("insert_llvm_intrinsic", llvmcall, return_type, input_types, args...)
where an LLVM vector type like<2 x i64>is translated to the Julia typeNTuple{2, VecElement{Int64}}
For example:
julia> const __m128i = NTuple{2, VecElement{Int64}};
julia> aesdec(a, roundkey) = ccall("llvm.x86.aesni.aesdec", llvmcall, __m128i, (__m128i, __m128i), a, roundkey);
julia> aesdec(__m128i((213132, 13131)), __m128i((31231, 43213)))
(VecElement{Int64}(-1627618977772868053), VecElement{Int64}(999044532936195731))
Automatic SIMD
The backend compiler for Julia is LLVM which can in some cases vectorize loops using the Loop Vectorizer
and it can even promote scalar code to SIMD operations using the SLP Vectorizer.
Automatic loop vectorization
Defining a simple loop that does an “axpy” like operation c .= a .* b
function axpy!(c::Array, a::Array, b::Array)
@assert length(a) == length(b) == length(c)
@inbounds for i in 1:length(a)
c[i] = a[i] * b[i]
end
end
and using the code introspection tool to inspect the generated LLVM IR, we see:
julia> V64 = Vector{Float64}
Array{Float64,1}
julia> code_llvm(axpy!, Tuple{V64, V64, V64})
...
%56 = fmul <4 x double> %wide.load, %wide.load24
%57 = fmul <4 x double> %wide.load21, %wide.load25
%58 = fmul <4 x double> %wide.load22, %wide.load26
%59 = fmul <4 x double> %wide.load23, %wide.load27
...
The type <4 x double> is in LLVM IR terminology a vector, and is here the resulting type
of adding two arguments of the same vector type.
As a note, LLVM has also decided to unroll the loop by a factor of four.
Moving on to look at the assembly, we see that indeed (at least on the computer of the author), this LLVM IR has turned into processor instructions that
does four multiplications in one instruction.
julia> code_native(axpy!, Tuple{V64, V64, V64})
vmulpd (%esi,%ecx,8), %ymm0, %ymm0
vmulpd 32(%esi,%ecx,8), %ymm1, %ymm1
vmulpd 64(%esi,%ecx,8), %ymm2, %ymm2
vmulpd 96(%esi,%ecx,8), %ymm3, %ymm3
The instruction vmulpd does “packed” double-precision floating point addition and
the ymm registers fit 256 bits which can thus fit four 64-bit floats.
Any type of control flow inside the loop will likely mean the loop will not vectorize. That is why @inbounds is important here,
otherwise we have control flow to the part that throws the bounds error.
Note that for reductions using non-associative arithmetic (like floating point airthmetic) you will have to tell the compiler that it is ok to reorder
the accumulations into the reduction variable using the @simd macro.
Automatic scalar vectorization
LLVM can also auto-vectorize scalar operations that follow a certain pattern. Here
is a function that does not contain any loop:
function mul_tuples(a::NTuple{4,Float64}, b::NTuple{4,Float64})
return (a[1]*b[1], a[2]*b[2], a[3]*b[3], a[4]*b[4])
end
The function mul_tuples just multiplies numbers from two tuples of length four and forms a new tuple.
The pattern here should be obvious, it is clear that the four additions
could be done at the same time.
LLVM can identify such patterns and generate code that uses SIMD. Again,
inspecting the code we find that SIMD instructions are used:
LLVM IR:
julia> code_llvm(mul_tuples, Tuple{NTuple{4,Float64}, NTuple{4, Float64}})
...
%7 = fmul <4 x double> %4, %6
...
Native code:
julia> code_native(mul_tuples, Tuple{NTuple{4,Float64}, NTuple{4, Float64}})
...
vmulpd (%edx), %ymm0, %ymm0
...
The scalar auto-vectorizer is quite impressive. It manages for example to very nicely vectorize a 4×4 matrix multiplication
in the StaticArrays.jl package which you can see doing something like
julia> using StaticArrays # import Pkg, Pkg.add("StaticArrays") to install
julia> @code_native rand(SMatrix{4,4}) * rand(SMatrix{4,4})
which almost only uses SIMD-instructions without StaticArrays.jl having to do any work for it.
SIMD using a vector library
While the auto-vectorizer can sometimes work pretty well, it quite easily gets confused. Alternatively, the data
is not laid out in such a way that it is allowed or beneficial to vectorize the code.
For example, trying a matrix multiplication of size
3×3 instead of 4×4 matrices in StaticArrays.jl and things are not so pretty anymore:
@code_native rand(SMatrix{3,3}) * rand(SMatrix{3,3})
vmovsd (%rdx), %xmm0 ## xmm0 = mem[0],zero
vmovsd 8(%rdx), %xmm7 ## xmm7 = mem[0],zero
vmovsd 16(%rdx), %xmm6 ## xmm6 = mem[0],zero
vmovsd 16(%rsi), %xmm11 ## xmm11 = mem[0],zero
vmovsd 40(%rsi), %xmm12 ## xmm12 = mem[0],zero
vmovsd 64(%rsi), %xmm9 ## xmm9 = mem[0],zero
vmovsd 24(%rdx), %xmm3 ## xmm3 = mem[0],zero
vmovupd (%rsi), %xmm4
vmovupd 8(%rsi), %xmm10
vmovhpd (%rsi), %xmm11, %xmm5 ## xmm5 = xmm11[0],mem[0]
vinsertf128 $1, %xmm5, %ymm4, %ymm5
vunpcklpd %xmm3, %xmm0, %xmm1 ## xmm1 = xmm0[0],xmm3[0]
vmovddup %xmm0, %xmm0 ## xmm0 = xmm0[0,0]
vinsertf128 $1, %xmm1, %ymm0, %ymm0
...
In the code above there is a lot of activity in the xmm registers (which are smaller than ymm).
indeed, if we benchmark the 3×3 matrix multiply we find that it is in fact slower than the 4×4 version (note that in the
benchmark below the matrix is wrapped in a Ref here to prevent the compiler from constant folding the benchmark loop):
julia> using BenchmarkTools # import Pkg; Pkg.add("BenchmarkTools") to install
julia> for n in (2,3,4)
s = Ref(rand(SMatrix{n,n}))
@btime $(s)[] * $(s)[]
end
2.711 ns (0 allocations: 0 bytes)
10.273 ns (0 allocations: 0 bytes)
6.059 ns (0 allocations: 0 bytes)
In these cases, we can explicitly vectorize the code using the SIMD vector library SIMD.jl.
SIMD.jl provides a type Vec{N, T} where N is the number of elements and T is the element type.
Vec{N, T} is similar to the LLVM <N x T> vector type and operations on Vec typically translate directly to LLVM operations:
For example, below we define some input data and a function g that do some simple arithmetic. We then look at the generated code.
julia> using SIMD
julia> a = Vec((1,2,3,4))
<4 x Int64>[1, 2, 3, 4]
julia> b = Vec((1,2,3,4))
<4 x Int64>[1, 2, 3, 4]
julia> g(a, b, c) = a * b + c;
julia> @code_llvm g(a, b, 3)
...
%res.i = mul <4 x i64> %7, %6
%8 = insertelement <4 x i64> undef, i64 %3, i32 0
%9 = shufflevector <4 x i64> %8, <4 x i64> undef, <4 x i32> zeroinitializer
%res.i1 = add <4 x i64> %res.i, %9
...
The mul <4 x i64> is the multiplication of the two vectors, and then the scalar 3 is “broadcasted” to a vector
and added to the result. Feel free to look at @code_native to see the native SIMD instructions.
We cand use SIMD.jl to write a faster 3×3 matrix multiplication:
function matmul3x3(a::SMatrix, b::SMatrix)
D1 = a.data; D2 = b.data
# Extract data from matrix into SIMD.jl Vec
SV11 = Vec((D1[1], D1[2], D1[3]))
SV12 = Vec((D1[4], D1[5], D1[6]))
SV13 = Vec((D1[7], D1[8], D1[9]))
# Form the columns of the resulting matrix
r1 = muladd(SV13, D2[3], muladd(SV12, D2[2], SV11 * D2[1]))
r2 = muladd(SV13, D2[6], muladd(SV12, D2[5], SV11 * D2[4]))
r3 = muladd(SV13, D2[9], muladd(SV12, D2[8], SV11 * D2[7]))
return SMatrix{3,3}((r1[1], r1[2], r1[3],
r2[1], r2[2], r2[3],
r3[1], r3[2], r3[3]))
end
Let’s compare this new version to the default the 3×3 version:
julia> s = Ref(rand(SMatrix{n,n}))
julia> @btime $(s)[] * $(s)[];
10.281 ns (0 allocations: 0 bytes)
julia> @btime matmul3x3($(s)[], $(s)[]);
4.392 ns (0 allocations: 0 bytes)
A guite significant improvement! The code for matmul3x3 could, of course, be generalized to work for more sizes, perhaps using a @generated function.
Using intrinsics
All we have done so far has been architecture independent.
If the CPU only supports an old version of SIMD or perhaps doesn’t support SIMD at all, LLVM will just compile the code
using the latest features that are available, falling back to scalar instructions if needed.
However, in some cases, we really do want to use a specific instruction in a certain instruction set supported by the CPU.
The idea for writing this blog post was from reading about
a new hashing library called “meowhash” was released.
It uses AES decryption which processors now has built-in instructions to perform.
Looking in the source code
we can see the macro:
#define Meow128_AESDEC(Prior, Xor) _mm_aesdec_si128((Prior), (Xor))
In the Intel Intrinsics Guide
this intrinsic is described as
Perform one round of an AES decryption flow on data (state) in a using the round key in
RoundKey, and store the result indst.
If we wanted to port meow-hash to Julia we would need to call this intrinsic in Julia, so how should we do that?
Firstly, Julia allows calling LLVM intrinsics through ccall.
We can for example call the pow intrinsic for two Float64 as:
julia> llvm_pow(a, b) = ccall("llvm.pow.f64", llvmcall, Float64, (Float64, Float64), a, b);
julia> llvm_pow(2.0, 3.0)
8.0
So, in order to call our AES decryption instruction, we need to know the corresponding LLVM intrinsic to _mm_aesdec_si128. Since Julia itself doesn’t provide us
with a way to get the intrinsic,
we need to ask a compiler that does. Fortunately, we can just ask Clang to emit the corresponding LLVM for us.
Using the Godbolt compiler webtool makes this very easy.
In the link to Godbolt we can see the following (slightly cleaned up):
define <2 x i64> @_Z6aesdecDv2_xS_(<2 x i64>, <2 x i64>) local_unnamed_addr #0 !dbg !262 {
%3 = call <2 x i64> @llvm.x86.aesni.aesdec(<2 x i64> %0, <2 x i64> %1) #3, !dbg !270
ret <2 x i64> %3, !dbg !271
}
So the intrinsic is called x86.aesni.aesdec. Now, we just need to know how to pass in the argument types which are
<2 x i64>. A normal Julia tuple of integers will not do because it gets passed to LLVM as an array
and not a vector
Instead, we need to send in a tuple with special elements of the type VecElement.
Julia treats a tuple of VecElements special and will pass it to LLVM as a vector.
All that is now left is to define some convenience typealias, create our inputs and call the intrinsic:
julia> const __m128i = NTuple{2, VecElement{Int64}};
julia> aesdec(a, roundkey) = ccall("llvm.x86.aesni.aesdec", llvmcall, __m128i, (__m128i, __m128i), a, roundkey);
julia> aesdec(__m128i((213132, 13131)), __m128i((31231, 43213)))
(VecElement{Int64}(-1627618977772868053), VecElement{Int64}(999044532936195731))
So now, we are in a position to port Meow Hash to Julia!
It should be stated that intrinsics should only be used as a last resort. It will lead to your code being less portable
and harder to maintain.
Conclusion
There are many ways of doing SIMD in Julia. From letting the compiler to do the the job to using a SIMD library, and finally
getting our hands dirty and use the intrinsics. Which way is best will depend on your application but hopefully, this helped a bit
with showing what options are available.