More docs and add Charms.Intrinsic

.github/workflows/elixir.yml

@@ -45,3 +45,9 @@ jobs:
       - name: Benchmark sort
         run: |
             mix run bench/sort_benchmark.exs
+      - name: Document
+        run: |
+            mix docs
+      - name: Package
+        run: |
+            mix archive.build

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,11 +1,31 @@
 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.
   """
 
   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}
@@ -39,13 +59,41 @@ 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_and_returns(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
+        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
   end
 end

lib/charms/defm.ex

@@ -51,47 +51,6 @@ 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
-  end
-
   @doc false
   def decompose_call_and_returns(call) do
     case call do
@@ -101,7 +60,7 @@ defmodule Charms.Defm do
   end
 
   @doc false
-  defp normalize_call(call) do
+  def normalize_call(call) do
     {name, args} = Macro.decompose_call(call)
 
     args =

lib/charms/env.ex

@@ -1,6 +1,7 @@
 defmodule Charms.Env do
-  use Beaver
+  use Charms.Intrinsic
 
+  @impl true
   def handle_intrinsic(:t, [], opts) do
     Beaver.ENIF.Type.env(opts)
   end

lib/charms/intrinsic.ex

@@ -0,0 +1,16 @@
+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()}]
+  @callback handle_intrinsic(atom(), [term()], opts()) :: term()
+
+  defmacro __using__(_) do
+    quote do
+      @behaviour Charms.Intrinsic
+      use Beaver
+    end
+  end
+end

lib/charms/pointer.ex

@@ -1,8 +1,9 @@
 defmodule Charms.Pointer do
-  use Beaver
+  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

lib/charms/prelude.ex

@@ -1,5 +1,5 @@
 defmodule Charms.Prelude do
-  use Beaver
+  use Charms.Intrinsic
   alias Beaver.MLIR.Dialect.{Arith, Func}
   @enif_functions Beaver.ENIF.functions()
   @binary_ops [:!=, :-, :+, :<, :>, :<=, :>=, :==, :&&, :*]
@@ -31,6 +31,7 @@ defmodule Charms.Prelude do
     v
   end
 
+  @impl true
   def handle_intrinsic(:result_at, [%MLIR.Value{} = v, i], _opts) when is_integer(i) do
     v
   end

lib/charms/simd.ex

@@ -1,8 +1,9 @@
 defmodule Charms.SIMD do
-  use Beaver
+  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

lib/charms/term.ex

@@ -1,6 +1,7 @@
 defmodule Charms.Term do
-  use Beaver
+  use Charms.Intrinsic
 
+  @impl true
   def handle_intrinsic(:t, [], opts) do
     Beaver.ENIF.Type.term(opts)
   end