Tricks to compile away step in calculation?

General Usage
6

In general

If you want the compiler to do something about this, you basically have to put the known argument (X) into the type domain. Then the compiler would be able to do constant propagation on X. However, this comes with significant caveats:

  1. Putting X into the type domain means putting it into a type parameter of a type. Only types, symbols, “plain data” (isbits) values and tuples thereof may be type parameters. In your case, this means that you can’t use Array, but you could use a static array, from the StaticArrays package. Use the isbits and isbitstype functions while developing to find out what is and what isn’t “plain data”.

  2. Getting the compiler involved means that you’ll have to pay the price of compilation and run time dispatch when constructing the closure. This may or may not be worth it, depending on how often you’ll run the closure after it’s constructed at great cost.

  3. For the Julia compiler to do constant propagation on a function (my_pinv in this case), it needs to know certain things about the function. If the compiler is not able to infer the necessary facts on it’s own, you may need to help it: Base.@assume_effects. While Base.@assume_effects is available in released Julia, IMO it’s significantly improved in the beta prerelease (v1.10), especially because there it’s possible to use Base.@assume_effects within the body of a function, which is safer. The function Base.infer_effects, which is not part of Julia’s public interface, must be used to find out what facts (“effects”) are inferred for a function.

  4. The Julia compiler doesn’t constant fold functions with @inbounds annotations. I hope a way around this will be made available.

These preconditions are tough to satisfy, but then, you’re definitely asking for much.

A simple example

This is an example that’s simpler than yours, but contrived.

# if necessary you can use a custom type instead of `Val`
to_value_domain(::Val{v}) where {v} = v

# this is instead of `ols` in your example
my_function(x, y) = (sin ∘ cos ∘ sin ∘ sin ∘ cos ∘ sin ∘ sin ∘ cos)(x) * y

# the usual way
function closed_function_without_run_time_dispatch(x)
  let x = x
    y -> my_function(x, y)
  end
end

# the "getting the compiler involved" way
function closed_function_with_run_time_dispatch(x)
  # warning: run time dispatch
  let x = Val(x)
    y -> my_function(to_value_domain(x), y)
  end
end

julia> f = closed_function_without_run_time_dispatch(0.3)
#1 (generic function with 1 method)

julia> g = closed_function_with_run_time_dispatch(0.3)
#3 (generic function with 1 method)

julia> f(0.1)
0.07238243561570594

julia> g(0.1)
0.07238243561570594

julia> @code_typed g(0.1)
CodeInfo(
1 ─ %1 = Base.mul_float(0.7238243561570593, y)::Float64
└──      return %1
) => Float64

julia> @code_typed f(0.1)
CodeInfo(
1 ─ %1 = Core.getfield(#self#, Symbol("#1#x"))::Float64
│   %2 = invoke Base.:(var"#_#103")($(QuoteNode(Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}()))::Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, sin ∘ cos ∘ sin ∘ sin ∘ cos ∘ sin ∘ sin ∘ cos::ComposedFunction{ComposedFunction{ComposedFunction{ComposedFunction{ComposedFunction{ComposedFunction{ComposedFunction{typeof(sin), typeof(cos)}, typeof(sin)}, typeof(sin)}, typeof(cos)}, typeof(sin)}, typeof(sin)}, typeof(cos)}, %1::Float64)::Float64
│   %3 = Base.mul_float(%2, y)::Float64
└──      return %3
) => Float64

Moving arrays into the type domain

As discussed above, Array is mutable, so not “plain data”. You could use something like this for moving arrays first into “plain data”, and then into the type domain:

using StaticArrays

function to_isbits(m::AbstractMatrix)
  s = size(m)

  # warning: run time dispatch
  SMatrix{s...}(m)
end

function to_type_domain(m::Matrix)
  # warning: run time dispatch
  Val{to_isbits(m)}()
end

Comparison with other languages

AFAIK Julia is currently the only programming language which allows the programmer to do something like I just described. In other languages you’d have to use metaprogramming, or templates or constexpr in C++, for example, while in Julia you can move values into the type domain. This is more powerful than other languages, because in other languages the precomputation would have to happen before compilation or during compilation, while in Julia the run time and compilation time are arbitrarily interleaved.

7 Likes