Charm Compiler Ideas

[jackalcooper]

Charm Compiler

The case for a Mojo for Elixir?

• Some developers still prefer a more functional style.
• Some developers dislike decorator/annotations.
• First class actor model, language-level support for preemptive scheduling.
• Open-source from day one, build in the public.
• Less legacy could also create more room for innovation, although there are less libraries to use in the beginning.

Compared to VM-based implementation of Erlang

• Conceptually, Charm should be a one-person project, which means avoiding introducing new semantic/paradigms don't exist in Erlang.
• Charm also has language-level support for preemptive scheduling. The compiler will keep track of function counters (very similar to what BEAM compiler has been doing now) and generate LLVM instructions to yield, instead of relying on VM at runtime.

New library or bindings could be enabled

• (godot extension](https://docs.godotengine.org/en/stable/tutorials/scripting/gdextension/what_is_gdextension.html)
• SIMD JSON, but implemented in Elixir
• Bootstrapped Erlang. To implement a variant of Erlang VM/Compiler in this Elixir DSL to achieve self-hosting at native level.

Reference

• Faster prefix sum computation with SIMD and Mojo
• Triton Vector Addition
• Let's Build a Baby Triton
• Port Triton to MLIR
• tinygo
• A Multi-Level Superoptimizer for Tensor Programs
• Triton recompiled to MLIR dialects
• Compile Triton to upstream MLIR
• Aten implementation in Triton

[josevalim]

Instead of building something that fully replaces the Erlang VM, I was thinking of an embedded language in Elixir that compiles to a NIF. The Mojo folks explain why they went with a separate language instead of an embedded DSL but I think Elixir will be both more powerful for embedding and the Erlang VM has a stronger and more isolated runtime to allow us to ship an embedded language. This also means the language can be a bit simpler (i.e. it doesn't have to worry about implementing concurrency).

Think of it as Cython (but not really) but compiling to MLIR. The Koka programming languages also has interesting ideas which we could explore if we decide to keep a more functional nature.

WDYT?

[jackalcooper]

embedded language in Elixir that compiles to a NIF.

I think this approach could be more practical and more useful for the community. There is one area I am not so sure about the NIF, which one is more ideal do you think:

1. export the final LLVM IR as objects and link them as an NIF shared library, the Erlang VM loads it.
2. there is a "generic" NIF wraps the LLVM JIT, and it can run the compiled functions of the embedded language.

[josevalim]

That's a very good question. Mojo mentions they use OrcJit but it is not immediately clear to me why. I guess they use it to perform both type and shape specializations at runtime? Similar to XLA?

[jackalcooper]

Based on my understanding and asking someone for help.
the orcjit here basically gives you a kind of flexibility. The changes of source code and the specializations for hardware require a lot of lazy compilation and linkages on demand.

[josevalim]

Makes sense. I assume you don't want to precompile all possible combinations for all possible types upfront. I guess in a way this is similar to what Julia does.

[jackalcooper]

I think the macro-based approach could really work. I am refactoring Beaver DSL into a macro-based one, whose form is closer to vanilla elixir. So when we have the C DSL, it would become a purer Elixir-to-Elixir AST transform, instead of worrying about MLIR CAPI calls. The C DSL's compiler's main jobs become clearer:

1. to figure out types and populate clauses like is(v2, Type.i(32))
2. to map functions in Elixir/Erlang to MLIR op calls

defm some_func2() do
  v0 = 0
  cond0 = true

  if cond0 do
    v1 = 0
    v1 + v1
  else
    v2 = 0
    v0 + v2
  end
end

Func.func some_func2(function_type: Type.function([], [Type.i(32)])) do
  region do
    block bb_entry() do
      v0 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) when is(v0, Type.i(32))
      cond0 = Arith.constant(true) when is(cond0, Type.i(32))

      cond_br(cond0) do
        bb1()
      else
        bb2(v0)
      end
    end

    block bb1() do
      v1 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) when is(v1, Type.i(32))
      add = Arith.addi(v0, v0) when is(add, Type.i(1))
      bb2(v1)
    end

    block bb2(arg) when is(arg, Type.i(32)) do
      v2 = Arith.constant(value: Attribute.integer(Type.i(32), 0)) when is(v2, Type.i(32))
      add = Arith.addi(arg, v2) when is(add, Type.i(32))
      Func.return(add)
    end
  end
end

[jackalcooper]

In this case it would be preprocessed by func

if it is processed by the surrounding macro like a func, it is what how Beaver works already. The only difference is that Beaver uses >>> instead of :: to "locate" supported constructs.

In other words, we will be able to provide Elixir dialects to each MLIR dialect that we want. Supporting macros will allow us to support complex constructs, such as for, conditionals, but in a granular and principled fashion. Anyone builds their own dialect. All you need is functions and macros.

The more I think about it, it sounds like what you want is a generic macro system expands Elixir to MLIR (or to the MLIR 1:1 expression in Elixir to be more specific). So that the job of building a compiler/dialect becomes composable by breaking the inference and transformations into these fine-grained rules (instead of explicitly doing recursive descent and walk). The major difference from Elixir's macro is that it is typed.

mlir_macro unless(expr, opts) do
  mlir_quote do
    if !mlir_unquote(expr), mlir_unquote(opts)
  end
end

mlir_macro if(condition, clauses) do
  build_if(condition, clauses) # generate MLIR ops
end

mlir_macro (arg0 :: i64) + (arg1 :: i64) do
  MLIR.arith.addi(arg0, arg1) :: i64 # because it is the target, no need to call mlir_quote
end

mlir_macro (arg0 :: i32) + (arg1 :: i32) do
  MLIR.arith.addi(arg0, arg1) :: i32
end

[josevalim]

Here is a proposal on how the syntax could look like:

MLIR."builtin.module"() do # single-block region
  var = MLIR."dialect.op1"[attribute_name: 42 :: i32]() :: {i1, i16, i32, i64}
  
  MLIR."dialect.op2"["other attribute": 42 :: i64]() do
    _ -> # unnamed block
      MLIR."dialect.innerop2"() :: {}
      MLIR."dialect.innerop3"(var._0, var._2, var._3) :: {}
    
    bb1(arg :: i32) ->
      MLIR."dialect.innerop4"() :: {}

    bb2(arg :: i64) ->
      MLIR."dialect.innerop4"() :: {}
  end
end

There are two annoyances:

• In Elixir, you can't express multiple return types without wrapping in a tuple. So we have to assume that (i32, i32 -> i32) is equivalent to (i32, i32 -> {i32}), the tuple wrapping for one element being optional.

• MLIR."dialect.op1"[attribute_name: 42 :: i32]() the access between the function name and the parens are not supported, but that's a relatively minor parser addition

We could also require the block names to be prefixed with ^, if we want to avoid ambiguity. We should not require SSA, I assume we can convert to SSA after the fact. Executing the code above should generate an Elixir AST with MLIR.Operand, MLIR.Block, MLIR.Region, but working on Elixir data structures rather than C. Then we do one pass that converts it to C/MLIR.

WDYT so far? I also need to do a pass over the type syntax.

[jackalcooper]

Generally I like this syntax very much! It matches and catches the essence of both MLIR and Elixir and indicates user to put operands/attributes where it supposed to be. I only have couple of questions in terms of usage and ergonomics. I call this "MLIR-Elixir-Format" as MEF for short

• For MEF code, within the body of a MLIR."builtin.module"() do, does it support coding in vanilla Elixir? I mean modules like Enum, String etc. The paradox here is that: if MEF supports mixing regular Elixir code, will the MEF compiler/transformer too completed because it needs to handle many edge cases. If MEF doesn't support mixing regular Elixir, why don't the user use MLIR's original textural format directly, considering they are so alike?
• Is the MLIR. in the beginning of an expression like MLIR."dialect.innerop3" an entry point for some extension mechanism? So that a compiler for MEF is a module called MLIR implementing a behavior? With that there will be things like XLA."stablehlo.negate"(op._0 :: tensor<2xi32>) :: tensor<2xi32> and EXLA can implement this behavior? Otherwise can we drop all the MLIR.s in the beginning because the parser can recognize the special format beginning with a string literal?
• What is the modularization and code reuse for MEF code? The most common case could be "everything is the same except for the type, so type should be a variable". It is OK if there is none if MEF is strictly an expression format instead of a "program" I guess.
• For more complicated type/attribute, how do they fit in the MEF. For instance, affine maps like #map3k = affine_map<(i, j) -> (i + 1, j*4*1 + 2)>. It is OK if we don't support them if we only want to cover enough MLIR functionality for a C level abstraction. More on this: https://mlir.llvm.org/docs/Dialects/Affine/#affine-maps

[josevalim]

I was able to sleep on this and refine it a bit more. We can probably use trailing keyword lists as the attributes syntax:

MLIR."builtin.module"() do # single-block region
  var = MLIR."dialect.op1"(attribute_name: 42 :: i32) :: {i1, i16, i32, i64}
  
  MLIR."dialect.op2"("other attribute": 42 :: i64) do
    _ -> # unnamed block
      MLIR."dialect.innerop2"() :: {}
      MLIR."dialect.innerop3"(var._0, var._2, var._3) :: {}
    
    bb1(arg :: i32) ->
      MLIR."dialect.innerop4"() :: {}

    bb2(arg :: i64) ->
      MLIR."dialect.innerop4"() :: {}
  end
end

This requires no new syntax and works today.

What I said about tuples in types may not be true:

So we have to assume that (i32, i32 -> i32) is equivalent to (i32, i32 -> {i32}), the tuple wrapping for one element being optional.

It depends if I can call var#1 in MLIR when no tuple types are returned. If I cannot, then they are distinct things (and better match Elixir semantics).

---

Now, to answer your questions:

For MEF code, within the body of a MLIR."builtin.module"() do, does it support coding in vanilla Elixir?

No, it is a subset of Elixir and it works pretty much like defn. Only what we allow is supported. When called at runtime, the func will return an AST (probably with an %MLIR.Operand{} struct at the root?).

in the beginning of an expression like MLIR."dialect.innerop3" an entry point for some extension mechanism?

No. func will find all calls to MLIR."dialect.innerop3"(...) and convert them to %MLIR.Operand{} structs. Unfortunately, we cannot drop MLIR. because we cannot call a string in Elixir. Even if we could, however, I would keep the MLIR suffix because we will add more features to this sublanguage. It will work like a onion:

• Layer 0 (core): the MLIR syntax above
• Layer 1: regular variables (expanded to SSA internally)
• Layer 2: support local/remote calls to other func (as in defn)
• Layer 3: support macros (as in defn)
• Layer ?: type parameters (for modularization?)
• Layer ?: ownership analysis

With these layers we will have a quite expressive syntax but all built in Elixir. The goal is to emit a pure-Elixir MLIR AST which will we can then compile to the C AST.

Elixir AST -> MLIR AST -> MLIR C

This means the core of the library can be written in pure Elixir.

#map3k = affine_map<(i, j) -> (i + 1, j41 + 2)>

As long as this translates to a MLIR C Operand or similar, we just need to come up with our own syntax to emit the same MLIR node.

[jackalcooper]

It depends if I can call var#1 in MLIR when no tuple types are returned. If I cannot, then they are distinct things (and better match Elixir semantics).

yes we can. the api is CAPI.mlirOperationGetResult(op, i)

[jackalcooper]

Following the expander POC snippet, refactor the prototype with Elixir 1.17's new functions in Macro module. It turns out to be exactly what I wanted but didn't know how to implement! There should be this mechanism to eagerly generate IR instead of evaluating Elixir AST, which was what I was doing and very error-prune.

Some pleasant changes after the refactor:

• no longer need to call a op macro when calling a function with no return value:

  # before
  def bar(a :: i32()) :: i32(), do: op func.return(a) :: []
  # now it is allowed
  def bar(a :: i32()) :: i32(), do: func.return(a)

  func.return here can be inferred as a MLIR operation.
• macro expansion is supported within a defm!

change doesn't land: remove the op macro

# i tried to remove the op macro
add_op = op llvm.add(a, b) :: i64()
sum = result_at(add_op, 0)
# which becomes
sum = llvm.add(a, b) :: i64()

::'s precedence is too low so it is expanded after sum = llvm.add(a, b) has been expanded. so an explicit op macro is still needed to process the [op] :: [type] expression as whole. A new value macro as syntax sugar is introduced.

sum = value llvm.add(a, b) :: i64()