Runtime macros or callsite transformations of generic functions.

[SimonDanisch]

I was thinking about codegen with different targets and optimization passes.
If we want to leave functions as generic as possible, this quite likely means we need to do the optimization at the callsite, not at function declaration.

So this:

@optimize_function_body function test(a,b,c)
    ...
end

Might be better like:

function test(a,b,c)
    ...
end
@optimize_function_body test(a,b,c) # obviously doesn't work with a macro
@fastmath test(a,b,c)

Since it leaves the function untouched. Also, the first version is not very composable, since you can't compose different optimization passes that easily.

This is why I propose callsite runtime macros.
I have only limited knowledge of codegen, but I actually implemented a very rough prototype, which offers the following interface:

# whereas stage can have the following values: lowered, typed, llvm
@callsite stage function optimization_pass(function_body::StageLambda, typed_arguments)
end
# alternatively
function optimization_pass(function_body, typed_arguments)
    body = code_typed_ast(function_body, typed_arguments)
end
# alternatively
function optimization_pass(function_body::Stage{:llvm}, typed_arguments)
end


_call(optimization_pass, function_to_optimize, (args,)) #_call is of course a staged function
# pass composition 
_call((body, args) -> pass3(pass2(pass1(body, args), args))), f, (2,3,4))
_call(rewrite_for_computeframework, f, (2,3,4)) # or spark!??

The simplest example would be callsite inlining.
More involved examples could be compilation of julia/llvm-ir for other targets, or flow analysis.

People already do similar things[1][2][3], so I'm wondering why we shouldn't offer a nice and stable interface for this!

Pros:

• keeps generic functions free of context specific optimizations
• will enable us to to implement fairly complex optimizations without burdening the core language
• executing generic code on different targets becomes simpler
• exposes one of the core features of Julia (great type inference, llvm-ir) to the crowd
• probably not worse than staged functions, right?

Cons:

• going further down the staged function rabbit hole
• could possibly recompile your whole program, while being very slow at that

I personally would like to have this in order to compile Julia functions for the GPU, without fiddling too much with Julia's internals.

CC: @ViralBShah, @carnaval, @vchuravy, @maleadt, @damiendr, @MikeInnes, @IntelLabs, @shashi

[JeffBezanson]

Could you describe more what the proposed feature actually does? Please explain step by step.

[SimonDanisch]

Sorry:

# Simple implementation, only working with the stage code_typed

# @generated function compiles a new version of `function_to_transform` for different transformations
@generated _call(transformation_func, function_to_transform, arguments...)
   # simplified version of https://github.com/SimonDanisch/Callsite.jl/blob/master/src/Callsite.jl#L29
   ast = code_typed(function_to_transform, arguments)::Vector{Expr}
   # should be transformation_func.instance, since the generated function only has access to the type
   expr_vector = transformation_func(ast , arguments)::Vector{Expr}
   quote 
        $(Expr(:block, expr_vector...))
   end
end
# probably better: args::NTuple{N, Pair{Symbol, DataType}}, to also pass arg name
function optimization_pass{N}(function_body::Vector{Expr}, args::NTuple{N, DataType}) 
     # transform the expression and return it
end
# applies `optimization_pass` to function body of f, compiles it as a function and calls the compiled 
# function with args...
_call(optimization_pass, f, (args...))

[edit]
Easiest example for a transformation function:

#callsite inlining

@callsite lowered function inline(body::Vector{Expr}, args)
    out = Any[Expr(:meta, :inline)]
    for elem in body
        elem != Expr(:meta, :noinline) && push!(out, elem)
    end
    return out
end

[vtjnash]

Implemented in #22440, although @jrevels has promised to continue adding more nice interfaces