More docs and add Charms.Intrinsic (#36)

.formatter.exs

@@ -3,7 +3,8 @@ locals_without_parens = [
   op: 1,
   value: 1,
   call: 1,
-  const: 1
+  const: 1,
+  defintrinsic: 1
 ]
 
 [

.github/workflows/elixir.yml

@@ -37,11 +37,12 @@ jobs:
       - name: Run tests
         run: mix test
       - name: Run dev build
-        run: |
-          mix
+        run: mix
+      - name: Document
+        run: mix docs
+      - name: Package
+        run: mix archive.build
       - name: Benchmark add
-        run: |
-            mix run bench/list_add_benchmark.exs
+        run: mix run bench/list_add_benchmark.exs
       - name: Benchmark sort
-        run: |
-            mix run bench/sort_benchmark.exs
+        run: mix run bench/sort_benchmark.exs

bench/vec_add_int_list.ex

@@ -1,4 +1,5 @@
 defmodule AddTwoIntVec do
+  @moduledoc false
   use Charms
   alias Charms.{SIMD, Term, Pointer}
 

lib/charms.ex

@@ -1,16 +1,34 @@
 defmodule Charms do
   @moduledoc """
   Documentation for `Charms`.
+
+  ## `defm` and intrinsic
+  There are two ways to define a function with `defm/2` or implement callbacks of `Charms.Intrinsic` behavior. The `defm/2` is a macro that generates a function definition in Charm. The intrinsic is a behavior that generates a function definition in MLIR.
+
+  The intrinsic is more flexible than `defm` because:
+  - Intrinsic can be variadic and its argument can be anything
+  - Intrinsic is suitable for the cases where directly writing or generating MLIR is more ideal
+  - An intrinsic should be responsible for its type check while the Charm’s type system is responsible for function call’s type check
+
+  The `defm` is more suitable for simple functions because it is designed to be as close to vanilla Elixir as possible. As a rule of thumb, use `defm` for simple functions and intrinsic for complex functions or higher-order(generic) function with type as argument.
+
+  ## `defm`'s differences from `Beaver.>>>/2` op expressions
+  - In `Beaver.>>>/2`, MLIR code are expected to mixed with regular Elixir code. While in `defm/2`, there is only Elixir code (a subset of Elixir, to be more precise).
+  - In `defm/2`, the extension of the compiler happens at the function level (define your intrinsics or `defm/2`s), while in `Beaver.>>>/2`, the extension happens at the op level (define your op expression).
+  - In `Beaver.>>>/2` the management of MLIR context and other resources are done by the user, while in `defm/2`, the management of resources are done by the `Charms` compiler.
+  - In `defm/2`, there is expected to be extra verifications built-in to the `Charms` compiler (both syntax and types), while in `Beaver.>>>/2`, there is none.
+
+  ## Caveats and limitations
+
+  - We need a explicit `call` in function call because the `::` special form has a parser priority  that is too low so a `call` macro is introduced to ensure proper scope.
+  - Being variadic, intrinsic must be called with the module name. `import` doesn't work with intrinsic functions while `alias` is supported.
   """
 
   defmacro __using__(opts) do
     quote do
-      import Charms.Defm
+      import Charms
       use Beaver
-      require Beaver.MLIR.Dialect.Func
-      alias Beaver.MLIR.Dialect.{Func, Arith, LLVM, CF}
-      alias Beaver.MLIR.{Type, Attribute}
-      import Type
+      import Beaver.MLIR.Type
 
       @before_compile Charms
       Module.register_attribute(__MODULE__, :defm, accumulate: true)
@@ -39,13 +57,42 @@ defmodule Charms do
     end
   end
 
-  def child_spec(mod, opts \\ [])
+  @doc """
+  define a function that can be JIT compiled
+  """
+  defmacro defm(call, body \\ []) do
+    {call, ret_types} = Charms.Defm.decompose_call_with_return_type(call)
+
+    call = Charms.Defm.normalize_call(call)
+    {name, args} = Macro.decompose_call(call)
 
-  def child_spec(mods, opts) when is_list(mods) do
-    %{id: Module.concat(mods), start: {Charms.JIT, :init, [mods, opts]}}
-  end
+    {:ok, env} =
+      __CALLER__ |> Macro.Env.define_import([], Charms.Defm, warn: false, only: :macros)
 
-  def child_spec(mod, opts) do
-    %{id: mod, start: {Charms.JIT, :init, [mod, opts]}}
+    [_enif_env | invoke_args] = args
+
+    invoke_args =
+      for {:"::", _, [a, _t]} <- invoke_args do
+        a
+      end
+
+    quote do
+      @defm unquote(Macro.escape({env, {call, ret_types, body}}))
+      def unquote(name)(unquote_splicing(invoke_args)) do
+        mfa = {unquote(env.module), unquote(name), unquote(invoke_args)}
+
+        cond do
+          @init_at_fun_call ->
+            {_, %Charms.JIT{engine: engine} = jit} = Charms.JIT.init(__MODULE__)
+            Charms.JIT.invoke(engine, mfa)
+
+          (engine = Charms.JIT.engine(__MODULE__)) != nil ->
+            Charms.JIT.invoke(engine, mfa)
+
+          true ->
+            &Charms.JIT.invoke(&1, mfa)
+        end
+      end
+    end
   end
 end

lib/charms/defm.ex

@@ -51,57 +51,17 @@ defmodule Charms.Defm do
   """
   defmacro cond_br(_condition, _clauses), do: :implemented_in_expander
 
-  @doc """
-  define a function that can be JIT compiled
-
-  ## Differences from `Beaver.>>>/2` op expressions
-  - In `Beaver.>>>/2`, MLIR code are expected to mixed with regular Elixir code. While in `defm/2`, there is only Elixir code (a subset of Elixir, to be more precise).
-  - In `defm/2`, the extension of the compiler happens at the function level (define your intrinsics or `defm/2`s), while in `Beaver.>>>/2`, the extension happens at the op level (define your op expression).
-  - In `Beaver.>>>/2` the management of MLIR context and other resources are done by the user, while in `defm/2`, the management of resources are done by the compiler.
-  - In `defm/2`, there is expected to be extra verifications built-in to the compiler (both syntax and types), while in `Beaver.>>>/2`, there is none.
-  """
-  defmacro defm(call, body \\ []) do
-    {call, ret_types} = decompose_call_and_returns(call)
-
-    call = normalize_call(call)
-    {name, args} = Macro.decompose_call(call)
-    env = __CALLER__
-    [_enif_env | invoke_args] = args
-
-    invoke_args =
-      for {:"::", _, [a, _t]} <- invoke_args do
-        a
-      end
-
-    quote do
-      @defm unquote(Macro.escape({env, {call, ret_types, body}}))
-      def unquote(name)(unquote_splicing(invoke_args)) do
-        if @init_at_fun_call do
-          {_, %Charms.JIT{}} = Charms.JIT.init(__MODULE__)
-        end
-
-        f =
-          &Charms.JIT.invoke(&1, {unquote(env.module), unquote(name), unquote(invoke_args)})
-
-        if engine = Charms.JIT.engine(__MODULE__) do
-          f.(engine)
-        else
-          f
-        end
-      end
-    end
+  @doc false
+  def decompose_call_with_return_type({:"::", _, [call, ret_type]}) do
+    {call, [ret_type]}
   end
 
-  @doc false
-  def decompose_call_and_returns(call) do
-    case call do
-      {:"::", _, [call, ret_type]} -> {call, [ret_type]}
-      call -> {call, []}
-    end
+  def decompose_call_with_return_type(call) do
+    {call, []}
   end
 
   @doc false
-  defp normalize_call(call) do
+  def normalize_call(call) do
     {name, args} = Macro.decompose_call(call)
 
     args =

lib/charms/defm/expander.ex

@@ -398,7 +398,7 @@ defmodule Charms.Defm.Expander do
     end
   end
 
-  @intrinsics Charms.Prelude.intrinsics()
+  @intrinsics Charms.Prelude.__intrinsics__()
   defp expand({fun, _meta, [left, right]}, state, env) when fun in @intrinsics do
     {left, state, env} = expand(left, state, env)
     {right, state, env} = expand(right, state, env)
@@ -477,7 +477,7 @@ defmodule Charms.Defm.Expander do
             Code.ensure_loaded(module)
 
             cond do
-              function_exported?(module, :handle_intrinsic, 3) ->
+              function_exported?(module, :__intrinsics__, 0) and fun in module.__intrinsics__() ->
                 {args, state, env} = expand(args, state, env)
 
                 {module.handle_intrinsic(fun, args, ctx: state.mlir.ctx, block: state.mlir.blk),
@@ -940,7 +940,7 @@ defmodule Charms.Defm.Expander do
   end
 
   defp expand_macro(_meta, Charms.Defm, :op, [call], _callback, state, env) do
-    {call, return_types} = Charms.Defm.decompose_call_and_returns(call)
+    {call, return_types} = Charms.Defm.decompose_call_with_return_type(call)
     {{dialect, _, _}, op, args} = Macro.decompose_call(call)
     op = "#{dialect}.#{op}"
     {args, state, env} = expand(args, state, env)
@@ -1055,7 +1055,6 @@ defmodule Charms.Defm.Expander do
 
   ## Helpers
 
-  @intrinsics Charms.Prelude.intrinsics()
   defp expand_remote(_meta, Kernel, fun, args, state, env) when fun in @intrinsics do
     {args, state, env} = expand(args, state, env)
 

lib/charms/env.ex

@@ -1,7 +1,13 @@
 defmodule Charms.Env do
-  use Beaver
+  @moduledoc """
+  Intrinsic module for BEAM environment's type.
+  """
+  use Charms.Intrinsic
 
+  @impl true
   def handle_intrinsic(:t, [], opts) do
     Beaver.ENIF.Type.env(opts)
   end
+
+  defintrinsic [:t]
 end

lib/charms/intrinsic.ex

@@ -0,0 +1,55 @@
+defmodule Charms.Intrinsic do
+  @moduledoc """
+  Behaviour to define intrinsic functions.
+  """
+  alias Beaver
+  @type opt :: {:ctx, MLIR.Context.t()} | {:block, MLIR.Block.t()}
+  @type opts :: [opt | {atom(), term()}]
+  @type ir_return :: MLIR.Value.t() | MLIR.Operation.t()
+  @type intrinsic_return :: ir_return() | (any() -> ir_return())
+  @doc """
+  Callback to implement an intrinsic.
+
+  Having different return types, there are two kinds of intrinsic functions:
+  - Regular: returns a MLIR value or operation.
+  - Higher-order: returns a function that returns a MLIR value or operation.
+
+  ## More on higher-order intrinsic
+  Higher-order intrinsic function can be variadic, which means it a list will be passed as arguments.
+  """
+  @callback handle_intrinsic(atom(), [term()], opts()) :: intrinsic_return()
+  Module.register_attribute(__MODULE__, :defintrinsic, accumulate: true)
+
+  @doc false
+  def collect_intrinsics(nil) do
+    raise ArgumentError, "no intrinsic functions defined"
+  end
+
+  def collect_intrinsics(attr_list) when length(attr_list) > 0 do
+    attr_list |> Enum.reverse() |> List.flatten() |> Enum.uniq()
+  end
+
+  defmacro __using__(_) do
+    quote do
+      @behaviour Charms.Intrinsic
+      use Beaver
+      @before_compile Charms.Intrinsic
+      import Charms.Intrinsic, only: :macros
+    end
+  end
+
+  defmacro defintrinsic(intrinsic_list) do
+    quote do
+      @defintrinsic unquote(intrinsic_list)
+    end
+  end
+
+  defmacro __before_compile__(_env) do
+    quote do
+      @defintrinsic_list @defintrinsic |> Charms.Intrinsic.collect_intrinsics()
+      def __intrinsics__() do
+        @defintrinsic_list
+      end
+    end
+  end
+end

lib/charms/jit.ex

@@ -147,11 +147,6 @@ defmodule Charms.JIT do
     beaver_raw_jit_invoke_with_terms(ref, to_string(Charms.Defm.mangling(mod, func)), args)
   end
 
-  def invoke(%MLIR.ExecutionEngine{} = engine, f, args)
-      when is_function(f, length(args)) do
-    apply(f, args).(engine)
-  end
-
   def destroy(module) do
     with %__MODULE__{ctx: ctx, engine: engine, owner: true} <-
            __MODULE__.LockedCache.get(module) do

lib/charms/pointer.ex

@@ -1,8 +1,12 @@
 defmodule Charms.Pointer do
-  use Beaver
+  @moduledoc """
+  Intrinsic module to work with pointers.
+  """
+  use Charms.Intrinsic
   alias Beaver.MLIR.{Type, Attribute}
   alias Beaver.MLIR.Dialect.{Arith, LLVM, Index}
 
+  @impl true
   def handle_intrinsic(:allocate, [elem_type], opts) do
     handle_intrinsic(:allocate, [elem_type, 1], opts)
   end
@@ -51,4 +55,6 @@ defmodule Charms.Pointer do
   def handle_intrinsic(:t, [], opts) do
     Beaver.Deferred.from_opts(opts, ~t{!llvm.ptr})
   end
+
+  defintrinsic [:t, :allocate, :load, :store, :element_ptr]
 end

lib/charms/prelude.ex

@@ -1,13 +1,12 @@
 defmodule Charms.Prelude do
-  use Beaver
+  @moduledoc """
+  Intrinsic module to define essential functions provided by Charms.
+  """
+  use Charms.Intrinsic
   alias Beaver.MLIR.Dialect.{Arith, Func}
   @enif_functions Beaver.ENIF.functions()
   @binary_ops [:!=, :-, :+, :<, :>, :<=, :>=, :==, :&&, :*]
 
-  def intrinsics() do
-    @enif_functions ++ [:result_at] ++ @binary_ops
-  end
-
   defp constant_of_same_type(i, v, opts) do
     mlir ctx: opts[:ctx], block: opts[:block] do
       t = MLIR.CAPI.mlirValueGetType(v)
@@ -31,10 +30,6 @@ defmodule Charms.Prelude do
     v
   end
 
-  def handle_intrinsic(:result_at, [%MLIR.Value{} = v, i], _opts) when is_integer(i) do
-    v
-  end
-
   def handle_intrinsic(:result_at, [l, i], _opts) when is_list(l) do
     l |> Enum.at(i)
   end
@@ -111,4 +106,6 @@ defmodule Charms.Prelude do
   def handle_intrinsic(_name, _args, _opts) do
     :not_handled
   end
+
+  defintrinsic @enif_functions ++ [:result_at] ++ @binary_ops
 end

lib/charms/simd.ex

@@ -1,8 +1,12 @@
 defmodule Charms.SIMD do
-  use Beaver
+  @moduledoc """
+  Intrinsic module for SIMD types.
+  """
+  use Charms.Intrinsic
   alias MLIR.Dialect.Arith
   alias MLIR.Type
 
+  @impl true
   def handle_intrinsic(:new, [type, width], opts) do
     fn literal_values ->
       mlir ctx: opts[:ctx], block: opts[:block] do
@@ -22,4 +26,6 @@ defmodule Charms.SIMD do
   def handle_intrinsic(:t, [type, width], _opts) do
     Type.vector([width], type)
   end
+
+  defintrinsic [:new, :t]
 end

lib/charms/term.ex

@@ -1,7 +1,13 @@
 defmodule Charms.Term do
-  use Beaver
+  @moduledoc """
+  Intrinsic module for SIMD type.
+  """
+  use Charms.Intrinsic
 
+  @impl true
   def handle_intrinsic(:t, [], opts) do
     Beaver.ENIF.Type.term(opts)
   end
+
+  defintrinsic [:t]
 end