Macro-based intrinsic (#48)

.formatter.exs

@@ -4,7 +4,6 @@ locals_without_parens = [
   value: 1,
   call: 1,
   const: 1,
-  defintrinsic: 1,
   defm: 2
 ]
 

bench/vec_add_int_list.ex

@@ -7,7 +7,7 @@ defmodule AddTwoIntVec do
     i_ptr = Pointer.allocate(i32())
     # TODO: remove the const here, when pointer's type can be inferred
     Pointer.store(const(0 :: i32()), i_ptr)
-    init = SIMD.new(i32(), 8).(0, 0, 0, 0, 0, 0, 0, 0)
+    init = SIMD.new(SIMD.t(i32(), 8), [0, 0, 0, 0, 0, 0, 0, 0])
 
     Enum.reduce(l, init, fn x, acc ->
       v_ptr = Pointer.allocate(i32())

lib/charms.ex

@@ -6,11 +6,12 @@ defmodule Charms do
   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
+  - It is possible for an intrinsic to return a MLIR type, while `defm` can only return value.
+  - Intrinsic function is always inline.
 
-  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.
+  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 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).
@@ -21,7 +22,6 @@ defmodule Charms do
   ## 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.
 
   ## Glossary of modules
 
@@ -36,6 +36,7 @@ defmodule Charms do
       import Charms
       use Beaver
       import Beaver.MLIR.Type
+      import Charms.Prelude
       @doc false
       def __use_ir__, do: nil
       @before_compile Charms

lib/charms/defm/expander.ex

@@ -364,7 +364,7 @@ defmodule Charms.Defm.Expander do
           MLIR.Block.add_args!(Beaver.Env.block(), arg_types, ctx: Beaver.Env.context())
 
           arg_values =
-            Range.new(0, length(args) - 1)
+            Range.new(0, length(args) - 1, 1)
             |> Enum.map(&MLIR.Block.get_arg!(Beaver.Env.block(), &1))
 
           state =
@@ -468,47 +468,62 @@ defmodule Charms.Defm.Expander do
     |> then(&{List.last(elem(&1, 0)), state, env})
   end
 
-  defp expand_intrinsics(loc, module, fun, args, state, env) do
+  defp expand_intrinsics(loc, module, intrinsic_impl, args, state, env) do
     {args, state, env} = expand(args, state, env)
     {params, state} = uniq_mlir_params(args, state)
 
-    case v =
-           module.handle_intrinsic(fun, params, args,
-             ctx: state.mlir.ctx,
-             block: state.mlir.blk,
-             loc: loc,
-             eval: fn ast ->
-               expand(
-                 ast,
-                 state,
-                 env
-               )
-             end,
-             params: params
-           ) do
+    v =
+      apply(module, intrinsic_impl, [
+        params,
+        %Charms.Intrinsic.Opts{
+          ctx: state.mlir.ctx,
+          args: args,
+          block: state.mlir.blk,
+          loc: loc,
+          eval: fn ast ->
+            expand(
+              ast,
+              state,
+              env
+            )
+          end
+        }
+      ])
+
+    case v do
       %m{} when m in [MLIR.Value, MLIR.Type, MLIR.Operation] ->
         {v, state, env}
 
       f when is_function(f) ->
         {f, state, env}
 
+      {:__block__, _, list} ->
+        # do not leak variables created in the macro
+        {v, _, _} = expand_list(list, state, env)
+
+        v
+        |> List.last()
+        |> then(&{&1, state, env})
+
       ast = {_, _, _} ->
-        {v, state, env} = expand(ast, state, env)
-        {List.last(v), state, env}
+        # do not leak variables created in the macro
+        {v, _state, _env} = expand(ast, state, env)
+        {v, state, env}
 
       other ->
         raise_compile_error(
           env,
-          "Unexpected return type from intrinsic #{module}.#{fun}: #{inspect(other)}"
+          "Unexpected return type from intrinsic #{module}.#{intrinsic_impl}: #{inspect(other)}"
         )
     end
   end
 
-  defp expand_magic_macros(loc, {module, fun, _arity} = mfa, args, state, env) do
+  defp expand_magic_macros(loc, {module, fun, arity} = mfa, args, state, env) do
     cond do
-      Code.ensure_loaded?(module) and function_exported?(module, :__intrinsics__, 0) and
-          fun in module.__intrinsics__() ->
-        expand_intrinsics(loc, module, fun, args, state, env)
+      Code.ensure_loaded?(module) and function_exported?(module, :__intrinsics__, 2) and
+          module.__intrinsics__(fun, arity) ->
+        intrinsic_impl = module.__intrinsics__(fun, arity)
+        expand_intrinsics(loc, module, intrinsic_impl, args, state, env)
 
       mfa == {Module, :__get_attribute__, 4} ->
         expand_get_attribute(args, state, env)
@@ -687,24 +702,32 @@ defmodule Charms.Defm.Expander do
     end
   end
 
-  @prelude_intrinsics Charms.Prelude.__intrinsics__()
-  defp expand({fun, _meta, [left, right]}, state, env) when fun in @prelude_intrinsics do
-    {left, state, env} = expand(left, state, env)
-    {right, state, env} = expand(right, state, env)
+  @intrinsics Charms.Kernel.macro_intrinsics() ++ Charms.Kernel.intrinsics()
+  defp expand({fun, _meta, args}, state, env) when fun in @intrinsics do
+    {args, state, env} = expand(args, state, env)
     loc = MLIR.Location.from_env(env)
-    {params, state} = uniq_mlir_params([left, right], state)
 
     try do
-      {Charms.Prelude.handle_intrinsic(fun, params, [left, right],
-         ctx: state.mlir.ctx,
-         block: state.mlir.blk,
-         loc: loc
-       ), state, env}
+      intrinsic_impl = Charms.Kernel.__intrinsics__(fun, length(args))
+      expand_intrinsics(loc, Charms.Kernel, intrinsic_impl, args, state, env)
+      |> then(fn {v, _, _} ->
+        if is_list(v) do
+          List.last(v)
+        else
+          v
+        end
+      end)
+      |> tap(fn v ->
+        unless match?(%MLIR.Value{}, v) do
+          raise_compile_error(env, "Expected a value, got: #{inspect(v)}")
+        end
+      end)
+      |> then(&{&1, state, env})
     rescue
       e ->
         raise_compile_error(
           env,
-          "Failed to expand prelude intrinsic #{fun}: #{Exception.message(e)}"
+          "Failed to expand kernel intrinsic #{fun}: #{Exception.message(e)}"
         )
     end
   end
@@ -1087,32 +1110,22 @@ defmodule Charms.Defm.Expander do
     {v, state, env}
   end
 
-  defp expand_macro(_meta, Kernel, :!, [value], _callback, state, env) do
-    {value, state, env} = expand(value, state, env)
-    type = MLIR.Value.type(value)
-    {value, state} = uniq_mlir_var(value, state)
-    {type, state} = uniq_mlir_var(type, state)
-
-    {not_value, state, env} =
-      quote do
-        one = const 1 :: unquote(type)
-        value arith.xori(unquote(value), one) :: unquote(type)
-      end
-      |> expand(state, env)
-
-    {List.last(not_value), state, env}
-  end
-
   defp expand_macro(_meta, Charms.Defm, :while, [expr, [do: body]], _callback, state, env) do
+    loc = MLIR.Location.from_env(env)
+
     v =
       mlir ctx: state.mlir.ctx, block: state.mlir.blk do
-        SCF.while [] do
+        SCF.while loc: loc do
           region do
             block _() do
               {condition, _state, _env} =
                 expand(expr, put_in(state.mlir.blk, Beaver.Env.block()), env)
 
-              SCF.condition(condition) >>> []
+              unless match?(%MLIR.Value{}, condition) do
+                raise_compile_error(env, "Expected a value, got: #{inspect(condition)}")
+              end
+
+              SCF.condition(condition, loc: loc) >>> []
             end
           end
 
@@ -1281,62 +1294,47 @@ defmodule Charms.Defm.Expander do
 
   ## Helpers
 
-  defp expand_remote(_meta, Kernel, fun, args, state, env) when fun in @prelude_intrinsics do
-    loc = MLIR.Location.from_env(env)
-    {args, state, env} = expand(args, state, env)
-    {params, state} = uniq_mlir_params(args, state)
-
-    {Charms.Prelude.handle_intrinsic(fun, params, args,
-       ctx: state.mlir.ctx,
-       block: state.mlir.blk,
-       loc: loc
-     ), state, env}
-  end
-
-  defp expand_remote(meta, module, fun, args, state, env) do
+  defp expand_remote(_meta, module, fun, args, state, env) do
     # A compiler may want to emit a :remote_function trace in here.
     state = update_in(state.remotes, &[{module, fun, length(args)} | &1])
     {args, state, env} = expand_list(args, state, env)
+    loc = MLIR.Location.from_env(env)
 
-    if module in [MLIR.Type] do
-      if fun in [:unranked_tensor, :complex, :vector] do
-        args
-      else
-        args ++ [[ctx: state.mlir.ctx]]
-      end
-      |> then(&{apply(module, fun, &1), state, env})
-    else
-      {{{:., meta, [module, fun]}, meta, args}, state, env}
+    cond do
+      Code.ensure_loaded?(module) and function_exported?(module, :__intrinsics__, 2) and
+          module.__intrinsics__(fun, length(args)) ->
+        intrinsic_impl = module.__intrinsics__(fun, length(args))
+        expand_intrinsics(loc, module, intrinsic_impl, args, state, env)
+
+      module in [MLIR.Type] ->
+        if fun in [:unranked_tensor, :complex, :vector] do
+          args
+        else
+          args ++ [[ctx: state.mlir.ctx]]
+        end
+        |> then(&{apply(module, fun, &1), state, env})
+
+      true ->
+        raise_compile_error(env, "function #{module}.#{fun}/#{length(args)} not found")
     end
   end
 
   defp expand_local(_meta, fun, args, state, env) do
     # A compiler may want to emit a :local_function trace in here.
     state = update_in(state.locals, &[{fun, length(args)} | &1])
     {args, state, env} = expand_list(args, state, env)
-    Code.ensure_loaded!(MLIR.Type)
-    loc = MLIR.Location.from_env(env)
 
-    if function_exported?(MLIR.Type, fun, 1) do
-      {apply(MLIR.Type, fun, [[ctx: state.mlir.ctx]]), state, env}
-    else
-      {params, state} = uniq_mlir_params(args, state)
-
-      case i =
-             Charms.Prelude.handle_intrinsic(fun, params, args,
-               ctx: state.mlir.ctx,
-               block: state.mlir.blk,
-               loc: loc
-             ) do
-        :not_handled ->
-          quote do
-            unquote(env.module).unquote(fun)(unquote_splicing(args))
-          end
-          |> expand_call_of_types([], state, env)
-
-        _ ->
-          {i, state, env}
+    try do
+      quote do
+        unquote(env.module).unquote(fun)(unquote_splicing(args))
       end
+      |> expand_call_of_types([], state, env)
+    rescue
+      e ->
+        raise_compile_error(
+          env,
+          "Failed to expand local function #{fun}/#{length(args)}: #{Exception.message(e)}"
+        )
     end
   end
 

lib/charms/env.ex

@@ -3,11 +3,9 @@ defmodule Charms.Env do
   Intrinsic module for BEAM environment's type.
   """
   use Charms.Intrinsic
+  alias Charms.Intrinsic.Opts
 
-  @impl true
-  def handle_intrinsic(:t, _params, [], opts) do
-    Beaver.ENIF.Type.env(opts)
+  defintrinsic t(), %Opts{ctx: ctx} do
+    Beaver.ENIF.Type.env(ctx: ctx)
   end
-
-  defintrinsic [:t]
 end