Fast inline compiler for numerical computing

Stream: lean4

[Tomas Skrivan]

As a hobby project, I would like to do some scientific/numerical computing in Lean. The main lure of Lean for me is the ability to formalize the mathematics behind these computations and hopefully use it to transform the code either for optimization or differentiation. (One motivating example is C++ library Eigen for linear algebra, that builds expressions involving matrices and vectors. Then at compile time it finds optimal strategy how to evaluate. )

However as an interpreted language, it would be too slow for numerical computing. Thus I'm thinking of writing a specialized compiler to speed up inner most loops of those computations that would not box types but keep them inlined as much as possible.

I do not attempt to create a compiler for arbitrary Lean code, but maybe I would just restrict to only a handful of types and compile non-dependent functions.

My hope is to inline types like:

- Float × Float to std::pair<double, double>
- Array (Float × Float) to std::vector<std::pair<double, double>> (std::vector is probably not a smart choice, but something like it with areference counter)
- Float^n to std::array<double, n>
i.e. mainly any fixed size vectors and matrices can be inlined/unboxed and arrays of these types will store them inline/unboxed.

Is this a good idea to attempt? Can I use the Lean.Compile.IR module in Lean for this? Or do I have to write my own? Any suggestions where should I start?

[Mario Carneiro]

However as an interpreted language, it would be too slow for numerical computing.

Nit: Lean 4 is not (exclusively or primarily) an interpreted language

[Mario Carneiro]

You will have a lot of trouble with that Float^n example because n is not normally a compile time constant in lean

[Mario Carneiro]

in fact there isn't even really a concept of compile time constants in the lean language so that kind of transformation will not be easy

[Mac]

Mario Carneiro said:

However as an interpreted language, it would be too slow for numerical computing.

Nit: Lean 4 is not (exclusively or primarily) an interpreted language

I would like to more strongly emphasis this point. Lean 4 is, for the perspective relevant to this question, a compiled language. Furthermore, Lean 4 is FAST. In fact, it is faster than many other functional programming languages and can be even be faster than other compiled languages in some cases. After all, a compiled Lean program is just a compiled C program with some, usually negligible, initialization .

To give an example, Lake (a build program for Lean written in Lean) is actually faster than using Make (a C program) for the same task.

[Mario Carneiro]

I think the aspect relevant to this question is that while lean is compiled, it has a (mostly) uniform data representation, which might incur performance costs in very high performance numerical computing applications. I'm not sure lean is ready to target that space, but it seems feasible to progressively expand the set of representations the compiler can handle

[Mac]

Yeah, I was about to say: if you are doing certain kinds of numerical computing, you may be looking to do some GPU accelerated things or use SIMD vectors. In which case, Lean does not yet support such things.

[Mac]

However, @Tomas Skrivan , your post seems to suggest you want to map Lean objects to C++ data structures like std::pair, std::vector, or std::array, which are no more efficient than their Lean equivalents.

[Mac]

In fact, the mappings you gave are more-or-less how those Lean types are defined (just in C rather than C++ terms).

[Tomas Skrivan]

Ok, fair point that about Lean and being interpreted language, well I do not even know formal definition of what interpreted language is :) Anyway, it is more or less irrelevant to what I'm asking.

Well one simple example of application I want to do is a simulation of particles. For this you need have an array of particles position, Array (Float × Float × Float), and at every time step you need to update these positions. I want really fast Array.map function that updates all particles.

Let's have a look at operations on Float × Float and what Lean actually compiles to.

Two simple functions, create a pair from two numbers and add two pairs together.

@[inline, export vec2_mk]
def vec2_mk (a b : Float) : Float × Float := (a, b)

@[inline, export vec2_add]
def vec2_add (a b : Float × Float) : Float × Float := (a.1 + b.1, a.2 + b.2)

The compiled code for vec2_mk:

  lean_object* vec2_mk(double x_1, double x_2) {
  _start:
    {
      lean_object* x_3; lean_object* x_4; lean_object* x_5;
      x_3 = lean_box_float(x_1);
      x_4 = lean_box_float(x_2);
      x_5 = lean_alloc_ctor(0, 2, 0);
      lean_ctor_set(x_5, 0, x_3);
      lean_ctor_set(x_5, 1, x_4);
      return x_5;
    }
  }

Both lean_box_float and lean_alloc_ctor allocate memory on the heap. So if you have a Array (Float × Float) you do not have a contiguous block of 2*sizeof(double)*array_length bytes but you have an array of pointers pointing all around the heap to pairs of doubles. (it is actually even worse, those doubles inside of a pair are boxed)

For numerical computing it is absolutely crucial to have data in a contiguous block of memory, as you need to utilize memory cache as much as possible. Quite often, you preload parts of a data from the heap to stack in a predictable manner and then you access it in what ever order you need.

The compile code for vec2_add

  lean_object* vec2_add(lean_object* x_1, lean_object* x_2) {
  _start:
    {
      lean_object* x_3; lean_object* x_4; double x_5; double x_6; double x_7; lean_object* x_8; lean_object* x_9; double x_10; double x_11; double x_12; lean_object* x_13; lean_object* x_14; lean_object* x_15;
      x_3 = lean_ctor_get(x_1, 0);
      lean_inc(x_3);
      x_4 = lean_ctor_get(x_2, 0);
      lean_inc(x_4);
      x_5 = lean_unbox_float(x_3);
      lean_dec(x_3);
      x_6 = lean_unbox_float(x_4);
      lean_dec(x_4);
      x_7 = x_5 + x_6;
      x_8 = lean_ctor_get(x_1, 1);
      lean_inc(x_8);
      lean_dec(x_1);
      x_9 = lean_ctor_get(x_2, 1);
      lean_inc(x_9);
      lean_dec(x_2);
      x_10 = lean_unbox_float(x_8);
      lean_dec(x_8);
      x_11 = lean_unbox_float(x_9);
      lean_dec(x_9);
      x_12 = x_10 + x_11;
      x_13 = lean_box_float(x_7);
      x_14 = lean_box_float(x_12);
      x_15 = lean_alloc_ctor(0, 2, 0);
      lean_ctor_set(x_15, 0, x_13);
      lean_ctor_set(x_15, 1, x_14);
      return x_15;
    }
  }

You have to first unbox all those floats, then you can add them and box them again. With all this you are touching heap and that is not acceptable for numerical computing.

I want these functions to compile to:

std::pair<double, double> vec2_mk(double a, double b)
{
   return std::pair{a, b};
}

std::pair<double, double> vec2_add(std::pair<double, double> a, std::pair<double, double> b)
{
   return std::pair{std::get<0>(a) + std::get<0>(b), std::get<1>(a) + std::get<1>(b)};
}

Actually, when you write a function Float → Float then it gets compiled nicely. An example:

@[inline, export foo]
def foo (a b : Float) : Float := ((Float.sin a) + (Float.cos b)) / (Float.exp a)

double foo(double x_1, double x_2) {
_start:
{
double x_3; double x_4; double x_5; double x_6; double x_7;
x_3 = sin(x_1);
x_4 = cos(x_2);
x_5 = x_3 + x_4;
x_6 = exp(x_1);
x_7 = x_5 / x_6;
return x_7;
}
}

Concerning Array T, there is a specialized implementation for Array UInt8 ~ lean_sarray_object and Array Char ~ lean_string_object. For every other type T, elements in the array are boxed. However, there is FloatArray that is an array of floats that are not boxed! Maybe I can get inspired by it and implement other arrays I need.

[Tomas Skrivan]

Mario Carneiro said:

You will have a lot of trouble with that Float^n example because n is not normally a compile time constant in lean

Yes this will be definitely a problem. But I can simply do #eval on an variable and see it it can be evaluated or not. Example:

def m := 2
#eval m -- 2

variable (n : Nat)
#eval n -- Error: (kernel) declaration has free variable `_eval`

Thus I might compile Float^n only if #eval n succeeds and use the value to turn it into std::array<double, n>.

[Leonardo de Moura]

@Tomas Skrivan You can get better code if you don't use polymorphic structures. Example:

structure FloatVec2 where
  a : Float
  b : Float

@[inline, export vec2_mk]
def vec2_mk (a b : Float) : FloatVec2 := { a, b }

@[inline, export vec2_add]
def vec2_add (a b : FloatVec2) : FloatVec2 := { a := a.1 + b.1, b := a.2 + b.2 }

We also have FloatArray. Which is an array of unboxed float values. However, the library for FloatArray is currently very small since nobody is using this kind of array.

[Tomas Skrivan]

Nice, this is probably because there is IRType.struct right? Maybe I can create FloatVec2Array, FloatVec3Array ... and if functions like get and set access element and component at the same time, I do not have to touch the heap.

[Mario Carneiro]

How far is lean from being able to do ad hoc monomorphization by using a macro to stamp out these array types and their libraries?

[Mario Carneiro]

My guess is that full C++/Rust style monomorphization will be out of scope for a long time

[Mario Carneiro]

Uniform data representation (aka pointers everywhere) is a necessity for languages that don't use monomorphization; just adding a bunch of additional data representations like FloatVec2Array without libraries and with an additional maintenance burden doesn't sound like a good idea

[Tomas Skrivan]

Makes sense. One thing I'm thinking of is implementing a small compiler with monomorphization(didn't know the term before) that is able to compile certain subset of Lean expressions and types. This would be used mainly for tight loops or when writing kernels for GPU.

[Mario Carneiro]

I think that writing your own compiler for this would be much more difficult than modifying / extending the lean compiler. Compilers have a lot of back end stuff that you would get for free if you modify an existing compiler.

[Mario Carneiro]

Well, that's not to say that the "extension" isn't just a way to hook up more low level semi-compiled lean expressions to the C output, and your actual compiler produces the semi-compiled lean expressions

[Mario Carneiro]

You might already have all the tools you need with enough calls to extern functions

[Mario Carneiro]

But if your compiler produces an infinite family of different ABI types, like array<double, n> or FloatVecNArray (for many values of n), then I don't see how lean.h is supposed to support that

[Tomas Skrivan]

I would expect that for each type you would have to also supply conversion to/from lean_object such those two types of codes can communicate.

[Mario Carneiro]

Put another way, I think the first thing to figure out is how to write lean code that has the performance characteristics you want, even if the code looks really ugly or low level. Then your compiler can be made to produce that ugly code from a more high level description

[Mario Carneiro]

Where would your types be defined?

[Mario Carneiro]

AFAIK lean code never defines new types, all types used in generated code are defined in lean.h

[Mac]

Mario Carneiro said:

My guess is that full C++/Rust style monomorphization will be out of scope for a long time

If you are talking about C++ templates (or Rust macros) -- they are both literally just macros. You could do everything they do (and more) with a macro In Lean.

[Mario Carneiro]

I was thinking not about rust macros, but rather rust generics

[Mario Carneiro]

rust macros are like lean macros, but rust generics are a whole type system thing that interfaces with the code generator

[Mario Carneiro]

you would be hard pressed to do all that with a macro

[Tomas Skrivan]

Mario Carneiro said:

Where would your types be defined?

My initial attempt is along these lines:

Define a C++ equivalent for each type and function:

class CppType (X : Type*) := ( typename : string )

-- define type for Prod
instance std_pair.CppType (A B : Type) [CppType A] [CppType B] : CppType (A×B) :=
{
  typename := "std::pair<" ++ CppType.typename A ++ ", " ++ CppType.typename B ++ ">"
}

class CppImpl {X : Type} (x : X) := { impl : string }

--- Define functions for Prod
variable {A B : Type} [CppType A] [CppType B]
instance  : CppImpl (@Prod.mk A B) := { impl := "std::make_pair" }
instance  : CppImpl (@Prod.fst A B) := { impl := "std::get<0>" }
instance  : CppImpl (@Prod.snd A B) := { impl := "std::get<1>" }

Once you define all the types you want to use and their functions, we can proceed with compiling an expression:

def pair_mk (a b : Float) : Float × Float := (a,b)

def compile : MetaM Unit := do
    let expr ← whnf (mkConst `pair_mk)
    ... Now you can traverse the expression and generate C++ code by gluing expression's c++ implementations. ...

This way I can generate some C++ code(still not working properly). I have no clue how to link it with the rest of the lean code.

The nice thing with this approach is that the C++ compiler will handle monomorphization for me.

[Christian Pehle]

Mario Carneiro said:

rust macros are like lean macros, but rust generics are a whole type system thing that interfaces with the code generator

Hobbes is an interesting example of a Haskell inspired language that does monomorphization for typeclasses https://github.com/morganstanley/hobbes, is that what you mean?

[Mario Carneiro]

Note that lean already has some monomorphization in the form of "specialization" of functions to certain types and typeclass instances, which sounds similar to hobbes

[Mario Carneiro]

but this is only an optimization, it doesn't (isn't supposed to) change observable behavior and it is not mandatory like C++/Rust monomorphization

[Christian Pehle]

Yes, I think what is mostly missing is unboxed structs and possibly some more base types like Float32.

[Mac]

Mac said:

If you are talking about C++ templates (or Rust macros) -- they are both literally just macros. You could do everything they do (and more) with a macro In Lean.

Just for fun, I wrote an example of how one could do def templates in Lean:

import Lean.Parser

open Lean Parser Command
open Term hiding ident

def declBinder :=
  simpleBinderWithoutType <|> bracketedBinder

def declBinders :=
  many (ppSpace >> declBinder)

@[commandParser]
def templateDefDecl := leading_parser
  "template " >> many1 Command.macroArg >> " def " >>
  incQuotDepth (ident >> declBinders >> (" := " <|> " where ") >> termParser)

def identAsStrLit (id : Syntax) : Syntax :=
  Syntax.mkStrLit (info := SourceInfo.fromRef id) <|
    id.getId.toString (escape := false)

@[macro templateDefDecl]
def expandTemplateDefDecl : Macro
| `(template $args* def $id $bs* := $x) =>
  let sym := identAsStrLit id
  `(macro $sym:strLit "[" $args* "]" : term => `(fun $[$bs]* => $x))
| _ => Macro.throwUnsupported

template t:term def add (a b : $t) := a + b
#check add[UInt64] 2 3 -- (fun a b => a + b) 2 3 : UInt64
#eval add[UInt64] 2 3 -- 5

[Mario Carneiro]

oh boy, I can't wait for SFINAE

[Mac]

Christian Pehle said:

Yes, I think what is mostly missing is unboxed structs and possibly some more base types like Float32.

Also, Lean already has unboxed doubles (that's what Float is) why do you particularly need an unboxed float?

[Mario Carneiro]

Why not? Lots of calculations don't need high precision floating point and would rather get the boost in throughput

[Mac]

@Mario Carneiro most systems are now 64-bit, thus 64-bit floating pointing arithmetic is just as fast (or faster) than 32-bit.

[Tomas Skrivan]

Also using float halfs the memory required compared to double.

[Mario Carneiro]

SIMD float instructions are 2x faster on 32 bit

[Mario Carneiro]

even if lean doesn't directly support SIMD, if the codegen is good enough you might be able to get autovectorization from the C compiler

[Mac]

@Mario Carneiro yes, the key part there being SIMD instructions which require SIMD vectors which is whole different can of worms.

[Mario Carneiro]

for instance, if you have a tight loop that doubles every element in a FloatArray, there is a high probability that will get autovectorized, and then you will be able to consume 2x more floats than doubles

[Tomas Skrivan]

Mario Carneiro said:

even if lean doesn't directly support SIMD, if the codegen is good enough you might be able to get autovectorization from the C compiler

Recently I was inspecting some generated assembly of C++ code I wrote. It had bunch of for loops for(int i=0;i<4;i++) and I got vectorization automatically.

[Mac]

Mario Carneiro said:

even if lean doesn't directly support SIMD, if the codegen is good enough you might be able to get autovectorization from the C compiler

iirc, auto-vectorization of floating point values is not on by default in most C compilers as that is another very different can of worms (as floating point rounding non-determinism then enters the picture).

[Mac]

however, I have never dug to deep into that ara, so I may be completely wrong there.

[Mario Carneiro]

https://www.godbolt.org/z/WGeqG835v seems to show that with a standard looking C++ loop you get autovectorization at -O3

[Mario Carneiro]

https://www.godbolt.org/z/49EczhxET <- also in C

[Mario Carneiro]

it is not clear whether an idiomatic lean loop will lower quite so beneficially as these C/C++ examples though. The most likely barrier to autovectorization of lean code is that all the jumps and function calls will not be inlined or inlined too late for the autovectorization recognizer to work

[Mac]

Welp, guess I was wrong. The fact that high optimization levels default to fast math surprises me.

[Mario Carneiro]

well, they say that -O3 is living on the wrong side of the guard rail

[Christian Pehle]

the purpose of undefined behaviour is to basically enable these optimisations ;).

[Christian Pehle]

https://godbolt.org/z/6vedK7WG3

[Mario Carneiro]

I don't think there are any fast-math style reassociations going on here though @Mac

[Mac]

Ah, now that I think about it, I may have been confusing a * b + c (add/mul or mul/add or whatever it is) with vectors.

[Mario Carneiro]

doing 8 adds in parallel is still IEEE conforming

[Mario Carneiro]

anyway, I think this is the main reason people care about float these days. The memory savings is nice but the double throughput is huge

[Mac]

@Mario Carneiro yep, that makes sense -- I was just under very mistaken assumptions.

[Mario Carneiro]

plus some double ops are actually slower

[Mac]

I think having a Float32 and Float64 split in Lean would be great.

[Mac]

Though I think Lean tends to avoid floating point numbers as they are very difficult to reason about, which was Lean's original goal.

[Mario Carneiro]

Indeed. I think Float32 isn't any worse than Float64 , but they are both pretty horrific for formal verification purposes

[Christian Pehle]

Mac said:

Ah, now that I think about it, I may have been confusing a * b + c (add/mul or mul/add or whatever it is) with vectors.

https://godbolt.org/z/Txee8dGe1 it does that too automatically.

[Mac]

@Christian Pehle Sorry, no, I meant the single add-mul or mul-add instruction (or something like that -- I can't remember which it was) not subsequent adds and muls on vectors. Scratch that, I am blind -- that is what is doing. :rolling_on_the_floor_laughing:

However, note it does not do that on plain -O3, but only with the necessary flag (which is what I was originally talking about).

[Mario Carneiro]

In particular @Tomas Skrivan , you should consider carefully whether you will actually get something that you can prove properties about at the end of all this. You might end up in no better position than if you had just written the code in C++, if everything is just axiomatized / opaque

[Mario Carneiro]

@Christian Pehle You did pass -mfma though, is that like a baby fast-math?

[Christian Pehle]

Mac said:

@Christian Pehle Sorry, no, I meant the single add-mul or mul-add instruction (or something like that -- I can't remember which it was) not subsequent adds and muls on vectors. Scratch that, I am blind -- that is what is doing. :rolling_on_the_floor_laughing:

However, note it does not do that on plain -O3, but only with the necessary flag (which is what I was originally talking about).

Ah yeah but that only depends on the architecture that the compiler assumes by default. If you use a gcc compiled to the target architecture it will do that on -O3 automatically (https://godbolt.org/z/3rKTaE74c)

[Mario Carneiro]

https://stackoverflow.com/a/15933677/890016 says:

The IEEE and C standards allow this when #pragma STDC FP_CONTRACT ON is in effect, and compilers are allowed to have it ON by default (but not all do). Gcc contracts into FMA by default (with the default -std=gnu*, but not -std=c*, e.g. -std=c++14). For Clang, it's only enabled with -ffp-contract=fast.

[Mac]

Sorry about derailing the discussion with my optimization confusions. Back on topic:

Tomas Skrivan said:

My initial attempt is along these lines:

Define a C++ equivalent for each type and function:

[...]

This way I can generate some C++ code(still not working properly). I have no clue how to link it with the rest of the lean code.

The nice thing with this approach is that the C++ compiler will handle monomorphization for me.

First, you should note that Lean outputs C, not C++. In fact, there is currently a big effort to drop C++ support entirely from Lean's compiler toolchain. Second, you can take a look at EmitC.lean to see how Lean currently goes about emitting C code.

[Mario Carneiro]

You could probably generate some cpp code and extern link to it from the lean-C code

[Tomas Skrivan]

Mario Carneiro said:

You could probably generate some cpp code and extern link to it from the lean-C code

Yes that is exactly what I'm thinking of doing. Maybe I just generate the C++ code and use only that in my final application without involving Lean's C code entirely.

[Tomas Skrivan]

Mario Carneiro said:

In particular Tomas Skrivan , you should consider carefully whether you will actually get something that you can prove properties about at the end of all this. You might end up in no better position than if you had just written the code in C++, if everything is just axiomatized / opaque

Well I'm will assume that I'm working with reals and hope for the best when actually running with floats. My main interest is to transform the code while respecting the mathematics. For example, I want to turn (A*B)*v to A*(B*v) to save computation for A,B matrices and v vector. On the level of floats (A*B)*v is not equal to A*(B*v), but I do not really care about that too much as long as I do not get some catastrophic rounding errors.
Doing these kind of optimizations in C++ is done with expression templates, for example the linear algebra library Eigen. I would like to take it a step further and do not only linear algebra but also calculus, differentiation etc. For a long time, I was trying doing it in C++, even switched to D for better templates. Now, I'm trying to do it in Lean and just generate the final C++ code.

One of my goals is to have code for finite elements where I can simply define energy as:

def energy (u : Ω → ℝ) : ℝ := ∫ ∥∇u∥² dx

That I can differentiate w.r.t. to u, assemble a finite dimensional system and find a minimizer.

[Mac]

Mario Carneiro said:

https://stackoverflow.com/a/15933677/890016 says:

The IEEE and C standards allow this when #pragma STDC FP_CONTRACT ON is in effect, and compilers are allowed to have it ON by default (but not all do). Gcc contracts into FMA by default (with the default -std=gnu*, but not -std=c*, e.g. -std=c++14). For Clang, it's only enabled with -ffp-contract=fast.

Note that this is for C (as can be seen here: https://godbolt.org/z/vrYzh39xn). On C++ (which @Christian Pehle was using) it will fuse mul and add (as can be seen here: https://godbolt.org/z/sEf33fccc).

[Mac]

Tomas Skrivan said:

Well I'm will assume that I'm working with reals and hope for the best when actually running with floats.

Unless I am mistaken (which given my current track record in this tread is completely possible), this is not how Lean works. You cannot reason between types like this in Lean.

[Mac]

Also, generally mathematical reasoning in Lean is non-computable so that is another hurdle with such an approach.

[Tomas Skrivan]

I just define an opaque type Real, define bunch of axioms defining them and when compiling Real will be turned to Float.

[Mac]

@Mario Carneiro, could that have soundness problems, because it feels like it could?

[Mario Carneiro]

Oh yes

[Mario Carneiro]

well I suppose it depends on the axioms but if you assume anything that isn't true about floats (and there are a lot of those) then you are on a very short path to unsoundness

[Tomas Skrivan]

But I'm not going to prove anything about what the program is actually doing in finite precision. I will be proving only what the program is doing in the idealized scenario of infinite precision.

[Mario Carneiro]

Actually I should reword that: As long as the axioms are consistent with each other (for example, they match what the real numbers do) then you won't be able to prove False in lean; but if they are not consistent with what floats do then you can get programs that don't compute what they are supposed to, and type violations and crashes (possibly including C UB)

[Tomas Skrivan]

Yes, I'm aware of that.

[Mario Carneiro]

in other words, you proved the correctness of a program that wasn't the one you ran

[Tomas Skrivan]

yes I know :)

[Tomas Skrivan]

My goal is not so much to prove correctness, but to optimize/transform/differentiate the code I write based on the underlining mathematics.

[Mario Carneiro]

Remind me why you can't do that in C++?

[Mario Carneiro]

I mean "Lean is nicer to work in" is a valid answer

[Mac]

@Mario Carneiro You'd have to do both separately, which I imagine would be a bit of a hassle. By doing both in Lean, you get a cheap representation of what the code you wrote was without have to remodel it in whatever proof system you are using.

[Tomas Skrivan]

To do differentiation, I need to work with some kind of expressions. Sure, I can implement my own system in C++ but lean does that for me already. Furthermore, differentiation is hard. I do not want to constrain myself to purely smooth functions or convex or what ever subset of functions.

[Mario Carneiro]

don't compilers have push button differentiation engines these days? Like enzyme

[Tomas Skrivan]

Like many differentiation packages just ignore that |x| is not differentiable at 0. Also I'm aware of only one language that can differentiate w.r.t. to a function.

[Christian Pehle]

in practice this is still an active area of research. there are a lot of tradeoffs and especially when it comes to optimisation in systems of ODEs and PDEs there are few/any compiler based approaches

[Christian Pehle]

zygotę and Julia come close.