Make entire function constant fold

Right now, in BetterExp, I have code that looks like

MAX_EXP(::Val{2},  ::Type{Float64}) =  1024.0                   # log2 2^1023*(2-2^-52)
MIN_EXP(::Val{2},  ::Type{Float64}) = -1075.0                   # log2 2^-1075
MAX_EXP(::Val{2},  ::Type{Float32}) =  128f0                    # log2 2^127*(2-2^-52)
MIN_EXP(::Val{2},  ::Type{Float32}) = -149f0                    # log2 2^-1075
MAX_EXP(::Val{ℯ},  ::Type{Float64}) =  709.7827128933845        # log 2^1023*(2-2^-52)
MIN_EXP(::Val{ℯ},  ::Type{Float64}) = -745.1332191019412076235  # log 2^-1075
MAX_EXP(::Val{ℯ},  ::Type{Float32}) =  88.72284f0               # log 2^127 *(2-2^-23)
MIN_EXP(::Val{ℯ},  ::Type{Float32}) = -103.97208f0              # log 2^-150
MAX_EXP(::Val{10}, ::Type{Float64}) =  308.25471555991675       # log10 2^1023*(2-2^-52)
MIN_EXP(::Val{10}, ::Type{Float64}) = -323.60724533877976       # log10 2^-1075
MAX_EXP(::Val{10}, ::Type{Float32}) =  38.53184f0               # log10 2^127 *(2-2^-23)
MIN_EXP(::Val{10}, ::Type{Float32}) = -45.1545f0                # log10 2^-150

I would much prefer to be able to write this as something like

f(::Val{N}, ::Type{Float64}) where {N} = 1023*log(N, 2) + log(N, 2-2^-52)

since it would be much cleaner. However, the compiler seems like it isn’t smart enough to do this constant folding automatically. Is there a way to force this function to compile down to a constant?

1 Like

If there’s only a short list of values, you could use @eval:

for N in (2, exp(1), 10)
    value = 1023*log(N, 2) + log(N, 2-2^-52)
    @eval f(::Val{$N}, ::Type{Float64}) = $value

This is precicely the intended usecase for @generated functions.

#+BEGIN_SRC julia
@generated f(::Val{N}, ::Type{Float64}) where {N} = 1023*log(N, 2) + log(N, 2-2^-52)

@code_typed f(Val(32), Float64)

1 ─     return 204.79999999999998
) => Float64

Be sure to read and understand though.

The essential thing is that you want to write them like

@generated function foo(x::T, y::U)
    # compute some things based on T and U, stuff here happens at compile time.
    # You can't use things like eval or Core.Compiler.return_type here.
        # An expression based on the results of the above computation.
        # This expr will get lowed to the runtime function body. 
        # If you put a value here rather than and Expr, your @generated function
        # be a constant function that just returns that value for any input of the 
        # correct types. 
        # The code that is returned here can't contain closures.

What are the pros and cons of the generated approach vs the macro approach? Obviously the generated version allows for arbitrary base arguments. Does it have a drawback?

@generated functions are for compile time computations when types are available. Macros are for parse-time computations where only syntax is available.

You can’t use compiler internals inside an @generated function and you can’t return a closure or certain other impure constructs. Note that the notion of the word ‘pure’ here is a little strange. For instance, array mutation is fine.

I’d argue that the @generated function approach is the correct one here because what you’re doing is really a type domain computation, not a syntax transformation. I’d encourage you to read the docs on them though, as I won’t be able to give you as accurate a picture of their limitations or structure.