A normal function outputs the result of the computation by the function. In contrast, a generated function outputs the code that implements the function. While generating this code, the generated function can only make use of the types of the arguments, not their values. In a sense, generated functions offer “on-demand code generation”. This mechanism is quite powerful and can be used when normal functions in combination with multiple dispatch cannot give you what you need.
To illustrate generated functions, we will build on the example of stack-allocated vectors from the previous post. We will extend our stack-allocated vector to a stack-allocated matrix, and we will use a generated function to implement matrix multiplication. Let’s start out by defining a stack-allocated vector and matrix.
struct StackMatrix{T,M,N,L}data::NTuple{L,T}endfunction StackVector(data::Vector{T})whereTreturnStackMatrix{T,length(data),1,length(data)}(Tuple(data))endfunction StackMatrix(data::Matrix{T})whereTM,N=size(data)returnStackMatrix{T,M,N,length(data)}(Tuple(data[:]))endThe type signature is StackMatrix{T, M, N, L} where T is the type of the elements of the matrix, M is the number of rows of the matrix, N is the number of columns of the matrix, and L = M * N is the total number of elements in the matrix;
even though L can always be computed from M and N, we need L in the type signature, because it specifies the length of the NTuple.
Before we implement multiplication of general StackMatrix{T, M, N, L}s, we first consider the case of StackMatrix{T, 2, 2, 4}s.
importBase:*function*(x::StackMatrix{T,2,2,4},y::StackMatrix{T,2,2,4})whereTx11,x21,x12,x22=x.datay11,y21,y12,y22=y.dataz11=x11*y11+x12*y21z21=x21*y11+x22*y21z12=x11*y12+x12*y22z22=x21*y12+x22*y22returnStackMatrix{T,2,2,4}((z11,z21,z12,z22))endLet’s check that the implementation is correct.
julia>x=randn(2,2);julia>y=randn(2,2);julia>x_stack=StackMatrix(x);julia>y_stack=StackMatrix(y);julia>x*y2×2Matrix{Float64}:-1.163610.8481590.355827-0.441428julia>reshape(collect((x_stack*y_stack).data),2,2)2×2Matrix{Float64}:-1.163610.8481590.355827-0.441428Nice! And it is quite a bit faster, too.
julia>@benchmark$x*$yBenchmarkTools.Trial:memoryestimate:112bytesallocsestimate:1--------------minimumtime:54.112ns(0.00%GC)mediantime:59.993ns(0.00%GC)meantime:62.529ns(1.27%GC)maximumtime:486.884ns(83.35%GC)--------------samples:10000evals/sample:973julia>@benchmark$(Ref(x_stack))[]*$(Ref(y_stack))[]BenchmarkTools.Trial:memoryestimate:0bytesallocsestimate:0--------------minimumtime:3.015ns(0.00%GC)mediantime:3.033ns(0.00%GC)meantime:3.077ns(0.00%GC)maximumtime:16.869ns(0.00%GC)--------------samples:10000evals/sample:1000The problem with multiplication of general StackMatrix{T, M, N, L}s is that the implementation depends on the particular values of M and N— for example, the variables z11, z21, et cetera.
We will use a generated function to automatically generate the implementation of the corresponding matrix multiplication.
This code-generation procedure depends on the values of M and N and will adapt the implementation accordingly.
Generated functions are defined with the macro @generated.
The implementation of multiplication of general StackMatrix{T, M, N, L}s as follows:
importBase:*@generatedfunction*(x::StackMatrix{T,K,M,L₁},y::StackMatrix{T,M,N,L₂})where{T,K,M,N,L₁,L₂}# Unpack `x`.tuple_x=Expr(:tuple,[Symbol("x_$(k)_$(m)")form=1:Mfork=1:K]...)unpack_x=:($tuple_x=x.data)# Unpack `y`.tuple_y=Expr(:tuple,[Symbol("y_$(m)_$(n)")forn=1:Nform=1:M]...)unpack_y=:($tuple_y=y.data)# Perform multiplication.mults=Vector{Expr}()fork=1:K,n=1:Nexpr=Expr(:call,:+,[:($(Symbol("x_$(k)_$(m)"))*$(Symbol("y_$(m)_$(n)")))form=1:M]...)push!(mults,:($(Symbol("z_$(k)_$(n)"))=$expr))end# Pack `z`.tuple_z=Expr(:tuple,[Symbol("z_$(k)_$(n)")forn=1:Nfork=1:K]...)pack_z=:(StackMatrix{T,K,N,L₃}($tuple_z))returnExpr(:block,unpack_x,unpack_y,mults...,:(L₃=K*N),:(return$pack_z))endIf we omit the macro @generated, we can call the implementation to inspect the generated code:
julia>x_stack*y_stackquote(x_1_1,x_2_1,x_1_2,x_2_2)=x.data(y_1_1,y_2_1,y_1_2,y_2_2)=y.dataz_1_1=x_1_1*y_1_1+x_1_2*y_2_1z_1_2=x_1_1*y_1_2+x_1_2*y_2_2z_2_1=x_2_1*y_1_1+x_2_2*y_2_1z_2_2=x_2_1*y_1_2+x_2_2*y_2_2L₃=K*NreturnStackMatrix{T,K,N,L₃}((z_1_1,z_2_1,z_1_2,z_2_2))endSweet! This looks very much like our earlier implementation of the two-by-two case. Let’s again check that the implementation is correct.
julia>x=randn(4,2);julia>y=randn(2,3);julia>x_stack=StackMatrix(x);julia>y_stack=StackMatrix(y);julia>x*y4×3Matrix{Float64}:0.125514-0.0135978-0.0283178-1.937560.4505591.173032.56769-0.365378-0.8459663.22549-0.602203-1.50065julia>reshape(collect((x_stack*y_stack).data),4,3)4×3Matrix{Float64}:0.125514-0.0135978-0.0283178-1.937560.4505591.173032.56769-0.365378-0.8459663.22549-0.602203-1.50065Like the two-by-two case, this implementation is quite a bit faster, too.
julia>@benchmark$x*$yBenchmarkTools.Trial:memoryestimate:176bytesallocsestimate:1--------------minimumtime:205.100ns(0.00%GC)mediantime:219.162ns(0.00%GC)meantime:229.711ns(0.74%GC)maximumtime:1.679μs(75.82%GC)--------------samples:10000evals/sample:530julia>@benchmark$(Ref(x_stack))[]*$(Ref(y_stack))[]BenchmarkTools.Trial:memoryestimate:0bytesallocsestimate:0--------------minimumtime:15.097ns(0.00%GC)mediantime:15.665ns(0.00%GC)meantime:16.987ns(0.00%GC)maximumtime:103.605ns(0.00%GC)--------------samples:10000evals/sample:997