The Design of Intrinsic

[jackalcooper]

RFC, the design of intrinsic

What is it?

Intrinsic is the feature in Charms that allows user to define a special kind of macro to generate defm AST but with access to MLIR CAPIs (provided by Beaver or other MLIR libraries in Elixir).
With that, Charms becomes more powerful and flexible to generate MLIR, on par with Mojo's meta-programming capability.

Use intrinsic to formalize Charms' defm

For a DSL, it is essential to define its syntax and semantics in a simple and comprehensive way.
Here is the rule of thumb of the atomic elements defm in Charms, or we can call it the leaf node of AST, basic node that no longer gets expanded:

• Only three kinds of elements in defm:
  • MLIR type: f32(), i32(), etc.
  • MLIR value: a operation's result or block argument
  • Elixir literal: integer, float, string, etc.
• Higher level elements are composed of these three kinds of elements, with macro as sugar, just like vanilla Elixir.

Here is the problem, with this design we manage to to keep the end-user experience consistent and simple, there is a gap between what defm itself and MLIR's basic entities. Besides type and value, MLIR has region, block, operation and attribute, etc. We need a way to bridge this gap, and intrinsic is the solution that emerged during the development of Charms.

Intrinsic is supposedly to be used by library authors or advanced users, who want to define a new kind of macro that can generate AST with access to MLIR CAPIs. It is a way to extend Charms' defm to cover more MLIR entities beyond its syntax.

Case study: Pointer.allocate/2

we want to define a new intrinsic Pointer.allocate/2 to allocate memory on stack, with the following signature:

# t can be a variable which is a MLIR type, and size a variable which is a MLIR value
Pointer.allocate(t, size)
# size can also be a integer literal
Pointer.allocate(i32(), 1)

if size is 1, omit it

Pointer.allocate(t)

Pointer.allocate/2 should be able to convert the size to i64 if it is not, and generate the AST for alloca operation.

And here are the two options to implement it:

Option 1, callback style

defmodule Charms.Pointer do
  # second argument params is the argument ast and third argument is the list of type, value or literal get passed
  def handle_intrinsic(:allocate, params, [_elem_type, size = %MLIR.Value{}], opts) do
    cast =
      cond do
        not MLIR.equal?(MLIR.Value.type(size), Type.i64(ctx: opts[:ctx])) ->
          quote do
            size = value arith.extsi(unquote(Enum.at(params, 1))) :: i64()
          end

        true ->
          size
      end

    quote do
      size = unquote(cast)
      value llvm.alloca(size, elem_type: unquote(Enum.at(params, 0))) :: "!llvm.ptr"
    end
  end
end

There should generate AST like this:

size = chv1059 # chv1059 is a uniq variable gets generated by Charms, means charms variable 1059
value llvm.alloca(size, elem_type: chv1027) :: "!llvm.ptr"

Option 2, def style

with this design, we have a second call in the definition, whose name must be params

defi allocate(elem_type, size), params(elem_type_var, size_var), opts do
  cast =
    cond do
      not MLIR.equal?(MLIR.Value.type(size), Type.i64(ctx: opts[:ctx])) ->
        quote do
          size = value arith.extsi(unquote(Enum.at(params, 1))) :: i64()
        end

      true ->
        size
    end

  quote do
    size = unquote(size_var)
    value llvm.alloca(size, elem_type: unquote(elem_type_var)) :: "!llvm.ptr"
  end
end

[polvalente]

I think that the first option is a bit too implicit. What I think would be nice is for a module that defines an intrinsic to have to declare it explicitly. The defi alternative does this, but the def handle_intrinsic alternative does not explicitly define an allocate function.

I assume the Beaver/Charms compiler is able to try and call an undefined function as an intrinsic, but ideally, you would at least have a "private name" like we have for defn: https://github.com/elixir-nx/nx/blob/9e2cd048de610151b85a27a183035bc0873fa77f/nx/lib/nx/defn/compiler.ex#L797

This way, intrinsics of the form Module.name(args) should always attempt to call either that name or Module.__intrinsic_#{name}__(args)

[polvalente]

Tangentially related, I feel like defintrinsic is a better name for clarity

[jackalcooper]

Tangentially related, I feel like defintrinsic is a better name for clarity

Agree

[jackalcooper]

you would at least have a "private name" like we have for defn

@polvalente I wanted to have this kind of function definition convention. But the problem in this case is that, there are extra arguments, like params (which is the ast) and at least an opts which contains MLIR block and context. My worry is that the convention will introduce inconsistency of arity (defintrinsic add(a,b) becomes def __intrinsic_add(a, b, params, opts)), and also might complicate things if the arguments have patten matching or \\ default. Does defn has measures or designs for these concerns?

[polvalente]

I see no issue with that. You can think of the public function as a shorthand for the private one, which also adds the extra params internally.

Defn does handle default args as expected too.

[jackalcooper]

I see no issue with that. You can think of the public function as a shorthand for the private one, which also adds the extra params internally.

In the private function, for instance Pointer.allocate(t) wants to call Pointer.allocate(t, 1), it is inevitable to introduce special macros like call_intrinsic right? I mean handle_intrinsic can call another handle_intrinsic because it is regular functions without wrapper.

[jackalcooper]

POC of new implementation

defmodule M do
  defmacro __using__(_) do
    quote do
      import M
      @before_compile M
      Module.register_attribute(__MODULE__, :intrinsic, accumulate: true)
    end
  end
  
  defmacro __before_compile__(_env) do
    quote bind_quoted: [] do
      @all_intrinsics @intrinsic |> Enum.reverse() |> List.flatten() |> Enum.uniq()
      
      for {name, _arity, intrinsic_name} <- @all_intrinsics do
        def __intrinsics__(unquote(name)) do
          unquote(intrinsic_name)
        end
      end
      
      def __intrinsics__() do
        @all_intrinsics
      end
    end
  end

  defmacro defintrinsic(call, opts \\ [], do: body) do
    {name, _meta, args} = call
    arity = length(args)
    # {opts, args} = List.pop_at(args, arity - 1)
    # arity = arity - 1
    
    intrinsic_name = :"__defintrinsic_#{name}__"
    
    opts_ast = 
      for {k, v} <- opts do
        quote do
          unquote(v) = opts[unquote(k)]
        end
      end
    
    quote do
      @intrinsic {unquote(name), unquote(arity), unquote(intrinsic_name)}
      def unquote(intrinsic_name)(unquote(args), opts) do
        unquote(opts_ast)
        unquote(body)
      end
      
      def unquote(name)(unquote_splicing(args)) do
        raise "Intrinsic #{unquote(name)}/#{unquote(arity)} cannot be called outside of a defm context"
      end
    end
    |> dbg()
  end
end

defmodule Pointer do
  use M
  
  defintrinsic allocate(elem_type), [a: a, b: b, nested: %{c: c}] do
    {a, b, c, elem_type, 123}
  end
  
  defintrinsic allocate(elem_type, length), [c: c] do
    {123, c}
  end
  
  defintrinsic load(out, elem_type) do
    123
  end
end

Pointer.__defintrinsic_allocate__([:type], [a: 1, b: 2, nested: %{c: 3}])
Pointer.__defintrinsic_allocate__([:type, 1], [c: 2])
Pointer.__defintrinsic_load__([:out, :type], [])

[jackalcooper]

• New syntax

@doc """
Allocates a single element of the given `elem_type`, returning a pointer to it.
"""
defintrinsic allocate(elem_type) do
  quote do
    Charms.Pointer.allocate(unquote(elem_type), 1)
  end
end

@doc """
Allocates an array of `size` elements of the given `elem_type`, returning a pointer to it.
"""
defintrinsic allocate(elem_type, size), %Opts{ctx: ctx, args: [_elem_type, size_v]} do
  cast =
    case size_v do
      i when is_integer(i) ->
        quote do
          const unquote(size_v) :: i64()
        end

      %MLIR.Value{} ->
        if MLIR.equal?(MLIR.Value.type(size_v), Type.i64(ctx: ctx)) do
          size
        else
          quote do
            value arith.extsi(unquote(size)) :: i64()
          end
        end
    end

  quote do
    size = unquote(cast)
    value llvm.alloca(size, elem_type: unquote(elem_type)) :: Pointer.t()
  end
end

• PR: Macro-based intrinsic #48

[jackalcooper]

New version

• Remove %Opts{} as argument. Replace it with __IR__ as variable
• Use bind_quoted to replace the concept of params
• PR: Use __IR__ in defintrinsic #49

@doc """
Allocates a single element of the given `elem_type`, returning a pointer to it.
"""
defintrinsic allocate(elem_type) do
  quote bind_quoted: [elem_type: elem_type] do
    Charms.Pointer.allocate(elem_type, 1)
  end
end

@doc """
Allocates an array of `size` elements of the given `elem_type`, returning a pointer to it.
"""
defintrinsic allocate(elem_type, size) do
  %Opts{ctx: ctx} = __IR__

  cast =
    case size do
      i when is_integer(i) ->
        quote bind_quoted: [size: i] do
          const size :: i64()
        end

      %MLIR.Value{} ->
        if MLIR.equal?(MLIR.Value.type(size), Type.i64(ctx: ctx)) do
          size
        else
          quote bind_quoted: [size: size] do
            value arith.extsi(size) :: i64()
          end
        end
    end

  quote bind_quoted: [elem_type: elem_type, size: cast] do
    value llvm.alloca(size, elem_type: elem_type) :: Pointer.t()
  end
end