Llvmcall with immediate values

I wanted to use some llvm instructions that require immediate values. Example:

using Base: llvmcall
const _u128 = NTuple{2, Base.VecElement{UInt64}}
u128(x,y)=(convert(Base.VecElement{UInt64},x),convert(Base.VecElement{UInt64},y))

@inline _aesni_keygenassist_4(a::_u128) =
    llvmcall(("declare <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64>, i8) nounwind readnone", 
        " %res = call <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64> %0, i8 8)
          ret <2 x i64> %res"), _u128, Tuple{_u128}, a) 

a=u128(1,2)
@show _aesni_keygenassist_4(a);
#_aesni_keygenassist_4(a) = (VecElement{UInt64}(0x6363636b63636363), VecElement{UInt64}(0x6363636b63636363))

Ok, so far so nice. This is ugly. Obvious attempt:

macro ll_keyassist(round) 
     code = " %res = call <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64> %0, i8 $(round))
     ret <2 x i64> %res"
     println("generating code: ", code)
               
     return :(key -> llvmcall(("declare <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64>, i8) nounwind readnone", code), _u128, Tuple{_u128}, key) )
end

This does not work:

julia> f(x)=(@ll_keyassist 8)(x)
generating code:  %res = call <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64> %0, i8 8)
     ret <2 x i64> %res
f (generic function with 1 method)

julia> f(a)
ERROR: error compiling f: error statically evaluating llvm IR argument

Is there a better way than defining _aesni_keygenassist_$(n) for all n?

Generated functions can be used here, but you need to encode the round number in a type.

julia> @generated function _aesni_keygenassist(::Type{Val{round}}, a::_u128) where {round}
           return quote
               llvmcall(("declare <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64>, i8) nounwind readnone", 
                         " %res = call <2 x i64> @llvm.x86.aesni.aeskeygenassist(<2 x i64> %0, i8 $(round))
                          ret <2 x i64> %res"), _u128, Tuple{_u128}, a) 
           end
       end
_aesni_keygenassist (generic function with 1 method)

julia> _aesni_keygenassist(Val{8}, a)
(VecElement{UInt64}(0x6363636b63636363), VecElement{UInt64}(0x6363636b63636363))

You may get some additional inspiration from the SIMD.jl package.

1 Like

Here’s a macro that does something like that:

https://github.com/tshort/ExportWebAssembly.jl/blob/e0380e066f376457af2664ebcc1b9ad16d1acea4/src/ExportWebAssembly.jl#L197-L209

For sufficiently complicated invocations, I’ve resorted to using LLVM.jl directly, like:

https://github.com/tshort/ExportWebAssembly.jl/blob/e0380e066f376457af2664ebcc1b9ad16d1acea4/lib/arrays/wasm-array.jl#L111-L128