diff --git a/.formatter.exs b/.formatter.exs index a08489d..6efb361 100644 --- a/.formatter.exs +++ b/.formatter.exs @@ -3,7 +3,8 @@ locals_without_parens = [ op: 1, value: 1, call: 1, - const: 1 + const: 1, + defintrinsic: 1 ] [ diff --git a/.github/workflows/elixir.yml b/.github/workflows/elixir.yml index 3e4e40b..dba2c86 100644 --- a/.github/workflows/elixir.yml +++ b/.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 diff --git a/bench/vec_add_int_list.ex b/bench/vec_add_int_list.ex index 0addeec..7f62907 100644 --- a/bench/vec_add_int_list.ex +++ b/bench/vec_add_int_list.ex @@ -1,4 +1,5 @@ defmodule AddTwoIntVec do + @moduledoc false use Charms alias Charms.{SIMD, Term, Pointer} diff --git a/lib/charms.ex b/lib/charms.ex index a4b0da5..9a000a9 100644 --- a/lib/charms.ex +++ b/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 diff --git a/lib/charms/defm.ex b/lib/charms/defm.ex index 7a010da..5a04214 100644 --- a/lib/charms/defm.ex +++ b/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 = diff --git a/lib/charms/defm/expander.ex b/lib/charms/defm/expander.ex index f4ef5d2..749ea1a 100644 --- a/lib/charms/defm/expander.ex +++ b/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) diff --git a/lib/charms/env.ex b/lib/charms/env.ex index cece8d7..218a4ad 100644 --- a/lib/charms/env.ex +++ b/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 diff --git a/lib/charms/intrinsic.ex b/lib/charms/intrinsic.ex new file mode 100644 index 0000000..5fc32f6 --- /dev/null +++ b/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 diff --git a/lib/charms/jit.ex b/lib/charms/jit.ex index 98215ec..7b51d48 100644 --- a/lib/charms/jit.ex +++ b/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 diff --git a/lib/charms/pointer.ex b/lib/charms/pointer.ex index 1fc8fb1..099fa67 100644 --- a/lib/charms/pointer.ex +++ b/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 diff --git a/lib/charms/prelude.ex b/lib/charms/prelude.ex index 8b349da..1692ccd 100644 --- a/lib/charms/prelude.ex +++ b/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 diff --git a/lib/charms/simd.ex b/lib/charms/simd.ex index ce2dada..f8c2624 100644 --- a/lib/charms/simd.ex +++ b/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 diff --git a/lib/charms/term.ex b/lib/charms/term.ex index 90ffe1c..9700df5 100644 --- a/lib/charms/term.ex +++ b/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 diff --git a/test/defm_test.exs b/test/defm_test.exs index fe2c177..ba4ce9d 100644 --- a/test/defm_test.exs +++ b/test/defm_test.exs @@ -34,10 +34,8 @@ defmodule DefmTest do {:ok, %Charms.JIT{}} = Charms.JIT.init(AddTwoInt, name: :add_int) engine = Charms.JIT.engine(:add_int) assert String.starts_with?(AddTwoInt.__ir__(), "ML\xefR") - assert Charms.JIT.invoke(engine, {AddTwoInt, :add, [1, 2, :arg_err]}) == 3 - assert Charms.JIT.invoke(engine, {AddTwoInt, :add, [1, "", :arg_err]}) == :arg_err - assert Charms.JIT.invoke(engine, &AddTwoInt.add/3, [1, 2, :arg_err]) == 3 - assert Charms.JIT.invoke(engine, &AddTwoInt.add/3, [1, "", :arg_err]) == :arg_err + assert AddTwoInt.add(1, 2, :arg_err).(engine) == 3 + assert AddTwoInt.add(1, "", :arg_err).(engine) == :arg_err assert :ok = Charms.JIT.destroy(:add_int) Charms.JIT.init(AddTwoInt)