From 2d2b2addbb54250c6f3c744b72b03136e8bcea54 Mon Sep 17 00:00:00 2001
From: Shenghang Tsai <jackalcooper@gmail.com>
Date: Wed, 27 Nov 2024 20:06:15 +0800
Subject: [PATCH] Macro-based intrinsic (#48)

---
 .formatter.exs              |   1 -
 bench/vec_add_int_list.ex   |   2 +-
 lib/charms.ex               |   7 +-
 lib/charms/defm/expander.ex | 188 ++++++++++++++++++------------------
 lib/charms/env.ex           |   8 +-
 lib/charms/intrinsic.ex     | 125 ++++++++++++++++++------
 lib/charms/kernel.ex        | 110 +++++++++++++++++++++
 lib/charms/pointer.ex       |  96 +++++++++---------
 lib/charms/prelude.ex       | 114 +++++++---------------
 lib/charms/simd.ex          |  44 ++++++---
 lib/charms/term.ex          |  13 +--
 test/const_test.exs         |   2 +-
 test/defm_test.exs          |  39 +++++++-
 test/expander_test.exs      |   4 +-
 test/vec_add_test.exs       |   4 +-
 15 files changed, 474 insertions(+), 283 deletions(-)
 create mode 100644 lib/charms/kernel.ex

diff --git a/.formatter.exs b/.formatter.exs
index bd3ccc0..f8e8fc4 100644
--- a/.formatter.exs
+++ b/.formatter.exs
@@ -4,7 +4,6 @@ locals_without_parens = [
   value: 1,
   call: 1,
   const: 1,
-  defintrinsic: 1,
   defm: 2
 ]
 
diff --git a/bench/vec_add_int_list.ex b/bench/vec_add_int_list.ex
index cc6de04..5f66b9a 100644
--- a/bench/vec_add_int_list.ex
+++ b/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())
diff --git a/lib/charms.ex b/lib/charms.ex
index a5920ad..41de02e 100644
--- a/lib/charms.ex
+++ b/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
diff --git a/lib/charms/defm/expander.ex b/lib/charms/defm/expander.ex
index 42a32b4..bf5df21 100644
--- a/lib/charms/defm/expander.ex
+++ b/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,32 +1294,28 @@ 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
 
@@ -1314,29 +1323,18 @@ defmodule Charms.Defm.Expander 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
 
diff --git a/lib/charms/env.ex b/lib/charms/env.ex
index 19bfce3..652b378 100644
--- a/lib/charms/env.ex
+++ b/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
diff --git a/lib/charms/intrinsic.ex b/lib/charms/intrinsic.ex
index 448ed07..fd7721a 100644
--- a/lib/charms/intrinsic.ex
+++ b/lib/charms/intrinsic.ex
@@ -1,4 +1,11 @@
 defmodule Charms.Intrinsic do
+  defmodule Opts do
+    @moduledoc """
+    Options for intrinsic functions.
+    """
+    defstruct [:ctx, :args, :block, :loc, :eval]
+  end
+
   @moduledoc """
   Behaviour to define intrinsic functions.
   """
@@ -7,49 +14,113 @@ defmodule Charms.Intrinsic do
   @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(), [Macro.t()], [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
+      Module.register_attribute(__MODULE__, :intrinsic, accumulate: true)
     end
   end
 
-  defmacro defintrinsic(intrinsic_list) do
+  defmacro defintrinsic(call, do: body) do
     quote do
-      @defintrinsic unquote(intrinsic_list)
+      defintrinsic(unquote(call), %Charms.Intrinsic.Opts{}, do: unquote(body))
     end
   end
 
-  defmacro __before_compile__(_env) do
+  defp unwrap_unquote(name) do
+    case name do
+      {:unquote, _, [name]} ->
+        name
+
+      _ ->
+        name
+    end
+  end
+
+  defp recompose_when_clauses(name, args, opts) do
+    intrinsic_name_ast =
+      {:unquote, [], [quote(do: :"__defintrinsic_#{unquote(unwrap_unquote(name))}__")]}
+
+    case opts do
+      {:when, when_meta, [opts | clauses]} ->
+        {:when, when_meta,
+         [
+           quote do
+             unquote(intrinsic_name_ast)(unquote(args), unquote(opts))
+           end
+           | clauses
+         ]}
+
+      _ ->
+        quote do
+          unquote(intrinsic_name_ast)(unquote(args), unquote(opts))
+        end
+    end
+  end
+
+  defp normalize_arg_names(args) do
+    for arg <- args do
+      case arg do
+        {arg_name, meta, nil} ->
+          arg_name
+          |> to_string()
+          |> String.trim_leading("_")
+          |> String.to_atom()
+          |> then(&{&1, meta, nil})
+
+        _ ->
+          arg
+      end
+    end
+  end
+
+  @doc """
+  To implement an intrinsic function
+  """
+  defmacro defintrinsic(call, opts, do: body) do
+    {name, _meta, args} = call
+    call = recompose_when_clauses(name, args, opts)
+    placeholder_args = normalize_arg_names(args)
+
+    # can't get the arity from length(args), because it might be an unquote_splicing, whose length is 1
+    placeholder =
+      quote generated: true do
+        def unquote(name)(unquote_splicing(placeholder_args)) do
+          arity = length([unquote_splicing(placeholder_args)])
+
+          raise "Intrinsic #{Exception.format_mfa(__MODULE__, unquote(name), arity)} cannot be called outside of a defm body"
+        end
+      end
+
     quote do
-      @defintrinsic_list @defintrinsic |> Charms.Intrinsic.collect_intrinsics()
-      def __intrinsics__() do
-        @defintrinsic_list
+      unquote(placeholder)
+      @doc false
+      def unquote(call) do
+        unquote(body)
       end
+
+      @intrinsic {unquote(unwrap_unquote(name)),
+                  :"__defintrinsic_#{unquote(unwrap_unquote(name))}__"}
+    end
+  end
+
+  defmacro __before_compile__(_env) do
+    quote bind_quoted: [] do
+      @all_intrinsics @intrinsic |> Enum.uniq()
+
+      for {name, intrinsic_name} <- @all_intrinsics do
+        def __intrinsics__(unquote(name), arity) do
+          if function_exported?(__MODULE__, unquote(name), arity) do
+            unquote(intrinsic_name)
+          end
+        end
+      end
+
+      def __intrinsics__(_, _), do: nil
     end
   end
 end
diff --git a/lib/charms/kernel.ex b/lib/charms/kernel.ex
new file mode 100644
index 0000000..26db1cf
--- /dev/null
+++ b/lib/charms/kernel.ex
@@ -0,0 +1,110 @@
+defmodule Charms.Kernel do
+  @moduledoc """
+  Intrinsic module to define functions in `Kernel`.
+  """
+  use Charms.Intrinsic
+  alias Charms.Intrinsic.Opts
+  alias Beaver.MLIR.Dialect.Arith
+  @binary_ops [:!=, :-, :+, :<, :>, :<=, :>=, :==, :*]
+  @unary_ops [:!]
+  @binary_macro_ops [:&&, :||]
+
+  defp constant_of_same_type(i, v, %Opts{ctx: ctx, block: block, loc: loc}) do
+    mlir ctx: ctx, block: block do
+      t = MLIR.CAPI.mlirValueGetType(v)
+
+      if MLIR.CAPI.mlirTypeIsAInteger(t) |> Beaver.Native.to_term() do
+        Arith.constant(value: Attribute.integer(t, i), loc: loc) >>> t
+      else
+        raise ArgumentError, "Not an integer type for constant, #{to_string(t)}"
+      end
+    end
+  end
+
+  @compare_ops [:!=, :==, :>, :>=, :<, :<=]
+  defp i_predicate(:!=), do: :ne
+  defp i_predicate(:==), do: :eq
+  defp i_predicate(:>), do: :sgt
+  defp i_predicate(:>=), do: :sge
+  defp i_predicate(:<), do: :slt
+  defp i_predicate(:<=), do: :sle
+
+  defp create_binary(op, operands, type, ctx, block, loc) do
+    mlir ctx: ctx, block: block do
+      case op do
+        op when op in @compare_ops ->
+          Arith.cmpi(operands, predicate: Arith.cmp_i_predicate(i_predicate(op)), loc: loc) >>>
+            Type.i1()
+
+        :- ->
+          Arith.subi(operands, loc: loc) >>> type
+
+        :+ ->
+          Arith.addi(operands, loc: loc) >>> type
+
+        :&& ->
+          Arith.andi(operands, loc: loc) >>> type
+
+        :|| ->
+          Arith.ori(operands, loc: loc) >>> type
+
+        :* ->
+          Arith.muli(operands, loc: loc) >>> type
+
+        _ ->
+          raise ArgumentError, "Unsupported operator: #{inspect(op)}"
+      end
+    end
+  end
+
+  for name <- @binary_ops ++ @binary_macro_ops do
+    defintrinsic unquote(name)(_left, _right),
+                 opts = %Opts{args: [left, right], ctx: ctx, block: block, loc: loc} do
+      {operands, type} =
+        case {left, right} do
+          {%MLIR.Value{} = v, i} when is_integer(i) ->
+            [v, constant_of_same_type(i, v, opts)]
+
+          {i, %MLIR.Value{} = v} when is_integer(i) ->
+            [constant_of_same_type(i, v, opts), v]
+
+          {%MLIR.Value{}, %MLIR.Value{}} ->
+            if not MLIR.equal?(MLIR.Value.type(left), MLIR.Value.type(right)) do
+              raise "args of binary op must be same type"
+            end
+
+            [left, right]
+
+          _ ->
+            raise ArgumentError,
+                  "Invalid arguments for binary operator: #{inspect(left)}, #{inspect(right)}"
+        end
+        |> then(fn [left, _] = operands -> {operands, MLIR.CAPI.mlirValueGetType(left)} end)
+
+      create_binary(unquote(name), operands, type, ctx, block, loc)
+    end
+  end
+
+  defintrinsic !_value, %Opts{args: [v]} do
+    t = MLIR.Value.type(v)
+
+    unless MLIR.CAPI.mlirTypeIsAInteger(t) |> Beaver.Native.to_term() do
+      raise ArgumentError, "Not an integer type to negate, unsupported type: #{to_string(t)}"
+    end
+
+    quote bind_quoted: [v: v, t: t] do
+      one = const 1 :: t
+      value arith.xori(v, one) :: t
+    end
+  end
+
+  @doc false
+  def intrinsics() do
+    @binary_ops
+  end
+
+  @doc false
+  def macro_intrinsics() do
+    @binary_macro_ops ++ @unary_ops
+  end
+end
diff --git a/lib/charms/pointer.ex b/lib/charms/pointer.ex
index 1f7147c..5b0ed10 100644
--- a/lib/charms/pointer.ex
+++ b/lib/charms/pointer.ex
@@ -4,42 +4,38 @@ defmodule Charms.Pointer do
   """
   alias Charms.Pointer
   use Charms.Intrinsic
-  alias Beaver.MLIR.{Type, Attribute}
-  alias Beaver.MLIR.Dialect.{Arith, LLVM, Index}
+  alias Charms.Intrinsic.Opts
+  alias Beaver.MLIR.{Type}
+  alias Beaver.MLIR.Dialect.{LLVM}
 
-  @impl true
-  def handle_intrinsic(:allocate, params, [elem_type], opts) do
-    handle_intrinsic(:allocate, params, [elem_type, 1], opts)
-  end
-
-  def handle_intrinsic(:allocate, _params, [elem_type, size], opts) when is_integer(size) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
-      size = Arith.constant(value: Attribute.integer(Type.i(32), size)) >>> ~t<i32>
-
-      size =
-        if MLIR.CAPI.mlirTypeIsAIndex(MLIR.Value.type(size)) |> Beaver.Native.to_term() do
-          Index.casts(size) >>> Type.i64()
-        else
-          size
-        end
-
-      LLVM.alloca(size, elem_type: elem_type) >>> ~t{!llvm.ptr}
+  @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
 
-  def handle_intrinsic(
-        :allocate,
-        [elem_type, size],
-        [_elem_type = %MLIR.Type{}, size_v = %MLIR.Value{}],
-        opts
-      ) do
+  @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 =
-      if MLIR.equal?(MLIR.Value.type(size_v), Type.i64(ctx: opts[:ctx])) do
-        size
-      else
-        quote do
-          size = value arith.extsi(unquote(size)) :: i64()
-        end
+      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
@@ -48,20 +44,33 @@ defmodule Charms.Pointer do
     end
   end
 
-  def handle_intrinsic(:load, _params, [type, ptr], opts) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
-      LLVM.load(ptr) >>> type
+  @doc """
+  Loads a value of `type` from the given pointer `ptr`.
+  """
+  defintrinsic load(type, ptr) do
+    quote do
+      value llvm.load(unquote(ptr)) :: unquote(type)
     end
   end
 
-  def handle_intrinsic(:store, _params, [val, ptr], opts) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
-      LLVM.store(val, ptr) >>> []
+  @doc """
+  Stores a value `val` at the given pointer `ptr`.
+  """
+  defintrinsic store(val, ptr) do
+    quote do
+      llvm.store(unquote(val), unquote(ptr))
     end
   end
 
-  def handle_intrinsic(:element_ptr, _params, [elem_type, ptr, n], opts) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
+  @doc """
+  Gets the element pointer of `elem_type` for the given base pointer `ptr` and index `n`.
+  """
+  defintrinsic element_ptr(_elem_type, _ptr, _n), %Opts{
+    ctx: ctx,
+    block: block,
+    args: [elem_type, ptr, n]
+  } do
+    mlir ctx: ctx, block: block do
       LLVM.getelementptr(ptr, n,
         elem_type: elem_type,
         rawConstantIndices: ~a{array<i32: -2147483648>}
@@ -69,9 +78,10 @@ defmodule Charms.Pointer do
     end
   end
 
-  def handle_intrinsic(:t, _params, [], opts) do
-    Beaver.Deferred.from_opts(opts, ~t{!llvm.ptr})
+  @doc """
+  Return the pointer type
+  """
+  defintrinsic t(), %Opts{ctx: ctx} do
+    Beaver.Deferred.create(~t{!llvm.ptr}, ctx)
   end
-
-  defintrinsic [:t, :allocate, :load, :store, :element_ptr]
 end
diff --git a/lib/charms/prelude.ex b/lib/charms/prelude.ex
index 33ad4c5..b4fbe17 100644
--- a/lib/charms/prelude.ex
+++ b/lib/charms/prelude.ex
@@ -3,24 +3,12 @@ defmodule Charms.Prelude do
   Intrinsic module to define essential functions provided by Charms.
   """
   use Charms.Intrinsic
+  alias Charms.Intrinsic.Opts
   alias Beaver.MLIR.Dialect.{Arith, Func}
   @enif_functions Beaver.ENIF.functions()
-  @binary_ops [:!=, :-, :+, :<, :>, :<=, :>=, :==, :&&, :||, :*]
 
-  defp constant_of_same_type(i, v, opts) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
-      t = MLIR.CAPI.mlirValueGetType(v)
-
-      if MLIR.CAPI.mlirTypeIsAInteger(t) |> Beaver.Native.to_term() do
-        Arith.constant(value: Attribute.integer(t, i)) >>> t
-      else
-        raise ArgumentError, "Not an integer type for constant, #{to_string(t)}"
-      end
-    end
-  end
-
-  defp wrap_arg({i, t}, opts) when is_integer(i) do
-    mlir ctx: opts[:ctx], block: opts[:block] do
+  defp wrap_arg({i, t}, %Opts{ctx: ctx, block: block}) when is_integer(i) do
+    mlir ctx: ctx, block: block do
       case i do
         %MLIR.Value{} ->
           i
@@ -39,82 +27,50 @@ defmodule Charms.Prelude do
     v
   end
 
-  @compare_ops [:!=, :==, :>, :>=, :<, :<=]
-  defp i_predicate(:!=), do: :ne
-  defp i_predicate(:==), do: :eq
-  defp i_predicate(:>), do: :sgt
-  defp i_predicate(:>=), do: :sge
-  defp i_predicate(:<), do: :slt
-  defp i_predicate(:<=), do: :sle
-
-  defp create_binary(op, operands, type, ctx, block) do
-    mlir ctx: ctx, block: block do
-      case op do
-        op when op in @compare_ops ->
-          Arith.cmpi(operands, predicate: Arith.cmp_i_predicate(i_predicate(op))) >>> Type.i1()
-
-        :- ->
-          Arith.subi(operands) >>> type
-
-        :+ ->
-          Arith.addi(operands) >>> type
-
-        :&& ->
-          Arith.andi(operands) >>> type
-
-        :|| ->
-          Arith.ori(operands) >>> type
+  defintrinsic result_at(_entity, _index),
+               %Opts{args: [%MLIR.Operation{} = op, i]} when is_integer(i) do
+    num_results = MLIR.CAPI.mlirOperationGetNumResults(op)
 
-        :* ->
-          Arith.muli(operands) >>> type
-      end
+    if i < num_results do
+      MLIR.CAPI.mlirOperationGetResult(op, i)
+    else
+      raise ArgumentError,
+            "Index #{i} is out of bounds for operation results, num results: #{num_results}"
     end
   end
 
-  def handle_intrinsic(:result_at, _params, [l, i], _opts) when is_list(l) do
-    l |> Enum.at(i)
-  end
-
-  def handle_intrinsic(:result_at, _params, [%MLIR.Operation{} = op, i], _opts) do
-    MLIR.CAPI.mlirOperationGetResult(op, i)
+  defintrinsic type_of(_value), %Opts{args: [v]} do
+    MLIR.Value.type(v)
   end
 
-  def handle_intrinsic(op, _params, [left, right], opts) when op in @binary_ops do
-    {operands, type} =
-      case {left, right} do
-        {%MLIR.Value{} = v, i} when is_integer(i) ->
-          [v, constant_of_same_type(i, v, opts)]
+  signature_ctx = MLIR.Context.create()
 
-        {i, %MLIR.Value{} = v} when is_integer(i) ->
-          [constant_of_same_type(i, v, opts), v]
+  for name <- @enif_functions do
+    {arg_types, _} = Beaver.ENIF.signature(signature_ctx, name)
+    args = Macro.generate_arguments(length(arg_types), __MODULE__)
 
-        {%MLIR.Value{}, %MLIR.Value{}} ->
-          [left, right]
-      end
-      |> then(fn [left, _] = operands -> {operands, MLIR.CAPI.mlirValueGetType(left)} end)
-
-    create_binary(op, operands, type, opts[:ctx], opts[:block])
-  end
+    defintrinsic unquote(name)(unquote_splicing(args)),
+                 opts = %Opts{args: args, ctx: ctx, block: block, loc: loc} do
+      {arg_types, ret_types} = Beaver.ENIF.signature(ctx, unquote(name))
+      args = args |> Enum.zip(arg_types) |> Enum.map(&wrap_arg(&1, opts))
 
-  def handle_intrinsic(name, _params, args, opts) when name in @enif_functions do
-    {arg_types, ret_types} = Beaver.ENIF.signature(opts[:ctx], name)
-    args = args |> Enum.zip(arg_types) |> Enum.map(&wrap_arg(&1, opts))
+      mlir ctx: ctx, block: block do
+        Func.call(args, callee: Attribute.flat_symbol_ref("#{unquote(name)}"), loc: loc) >>>
+          case ret_types do
+            [ret] ->
+              ret
 
-    mlir ctx: opts[:ctx], block: opts[:block] do
-      Func.call(args, callee: Attribute.flat_symbol_ref("#{name}"), loc: opts[:loc]) >>>
-        case ret_types do
-          [ret] ->
-            ret
-
-          [] ->
-            []
-        end
+            [] ->
+              []
+          end
+      end
     end
   end
 
-  def handle_intrinsic(_name, _params, _args, _opts) do
-    :not_handled
-  end
+  MLIR.Context.destroy(signature_ctx)
 
-  defintrinsic @enif_functions ++ [:result_at] ++ @binary_ops
+  @doc false
+  def intrinsics() do
+    @enif_functions ++ [:result_at]
+  end
 end
diff --git a/lib/charms/simd.ex b/lib/charms/simd.ex
index 692998e..04833f6 100644
--- a/lib/charms/simd.ex
+++ b/lib/charms/simd.ex
@@ -5,27 +5,43 @@ defmodule Charms.SIMD do
   use Charms.Intrinsic
   alias MLIR.Dialect.Arith
   alias MLIR.Type
+  alias Charms.Intrinsic.Opts
 
-  @impl true
-  def handle_intrinsic(:new, params, [type, width], opts) do
-    fn literal_values ->
-      mlir ctx: opts[:ctx], block: opts[:block] do
-        values = Enum.map(literal_values, &Attribute.integer(type, &1))
+  @doc """
+  Return the constant value of the given `type` and `literal_values`
+  """
+  defintrinsic new(_type, _literal_values), %Opts{
+    args: [type, literal_values],
+    ctx: ctx,
+    block: block
+  } do
+    mlir ctx: ctx, block: block do
+      element_type = MLIR.CAPI.mlirShapedTypeGetElementType(type)
+
+      if MLIR.is_null(element_type) do
+        raise "element type is null"
+      end
+
+      width = MLIR.CAPI.mlirShapedTypeGetDimSize(type, 0) |> Beaver.Native.to_term()
 
-        if Enum.count(values) != width do
-          raise ArgumentError, "expected #{width} values, got #{length(values)}"
-        end
+      if Enum.count(literal_values) != width do
+        raise ArgumentError, "expected #{width} values, got #{length(literal_values)}"
+      end
 
-        t = handle_intrinsic(:t, params, [type, width], opts)
-        value = Attribute.dense_elements(values, t, opts)
-        Arith.constant(value: value) >>> t
+      if width <= 0 do
+        raise ArgumentError, "width must be a positive integer"
       end
+
+      values = Enum.map(literal_values, &Attribute.integer(element_type, &1))
+      value = Attribute.dense_elements(values, type, ctx: ctx)
+      Arith.constant(value: value) >>> type
     end
   end
 
-  def handle_intrinsic(:t, _params, [type, width], _opts) do
+  @doc """
+  Return the vector type of the given `type` and `width`
+  """
+  defintrinsic t(_type, _width), %Opts{args: [type, width]} do
     Type.vector([width], type)
   end
-
-  defintrinsic [:new, :t]
 end
diff --git a/lib/charms/term.ex b/lib/charms/term.ex
index 88f80b5..d99f722 100644
--- a/lib/charms/term.ex
+++ b/lib/charms/term.ex
@@ -1,13 +1,14 @@
 defmodule Charms.Term do
   @moduledoc """
-  Intrinsic module for SIMD type.
+  Intrinsic module for Erlang term type.
   """
   use Charms.Intrinsic
+  alias Charms.Intrinsic.Opts
 
-  @impl true
-  def handle_intrinsic(:t, _params, [], opts) do
-    Beaver.ENIF.Type.term(opts)
+  @doc """
+  Return the Erlang term type.
+  """
+  defintrinsic t(), %Opts{ctx: ctx} do
+    Beaver.ENIF.Type.term(ctx: ctx)
   end
-
-  defintrinsic [:t]
 end
diff --git a/test/const_test.exs b/test/const_test.exs
index 9e4f324..488f993 100644
--- a/test/const_test.exs
+++ b/test/const_test.exs
@@ -16,7 +16,7 @@ defmodule ConstTest do
     end
 
     assert_raise CompileError,
-                 "test/const_test.exs:13: Unsupported type for const macro: tensor<*xf64>",
+                 ~r"test/const_test.exs:13: Unsupported type for const macro: tensor<\*xf64>",
                  f
   end
 end
diff --git a/test/defm_test.exs b/test/defm_test.exs
index ba7391e..03e0cc1 100644
--- a/test/defm_test.exs
+++ b/test/defm_test.exs
@@ -49,7 +49,7 @@ defmodule DefmTest do
 
   test "invalid return of absent alias" do
     assert_raise CompileError,
-                 "test/defm_test.exs:#{__ENV__.line + 5}: invalid return type",
+                 ~r"test/defm_test.exs:#{__ENV__.line + 5}: invalid return type",
                  fn ->
                    defmodule InvalidRet do
                      use Charms
@@ -63,13 +63,13 @@ defmodule DefmTest do
 
   test "invalid arg of absent alias" do
     assert_raise CompileError,
-                 "test/defm_test.exs:#{__ENV__.line + 6}: invalid argument type #2",
+                 ~r"test/defm_test.exs:#{__ENV__.line + 6}: invalid argument type #2",
                  fn ->
-                   defmodule InvalidRet do
+                   defmodule InvalidArgType do
                      use Charms
                      alias Charms.Term
 
-                     defm my_function(env, arg1 :: Pointer.t(), arg2) :: Term.t() do
+                     defm my_function(env, arg1 :: Invalid.t(), arg2) :: Term.t() do
                        func.return(arg2)
                      end
                    end
@@ -172,4 +172,35 @@ defmodule DefmTest do
       end
     end
   end
+
+  test "negate" do
+    assert_raise CompileError, ~r/Not an integer type to negate, unsupported type: f32/, fn ->
+      defmodule NegateFloatType do
+        use Charms
+        alias Charms.Term
+
+        defm foo() do
+          zero = const 0.0 :: f32()
+          !zero
+        end
+      end
+    end
+  end
+
+  test "type of" do
+    defmodule TypeOf do
+      use Charms
+      alias Charms.{Term, Pointer}
+
+      defm foo(env, a) :: Term.t() do
+        b = const 1 :: i32()
+        i_ptr = Pointer.allocate(type_of(b))
+        enif_get_int(env, a, i_ptr)
+        sum = Pointer.load(type_of(b), i_ptr) + b
+        enif_make_int(env, sum)
+      end
+    end
+
+    assert TypeOf.foo(2) == 3
+  end
 end
diff --git a/test/expander_test.exs b/test/expander_test.exs
index 5a37099..0575180 100644
--- a/test/expander_test.exs
+++ b/test/expander_test.exs
@@ -166,7 +166,7 @@ defmodule POCTest do
 
     test "op not found" do
       assert_raise CompileError,
-                   "example.exs: Unknown MLIR operation to create: cf.ar, did you mean: cf.br",
+                   ~r"example.exs: Unknown MLIR operation to create: cf.ar, did you mean: cf.br",
                    fn ->
                      quote do
                        defmodule ReturnPassedArg do
@@ -181,7 +181,7 @@ defmodule POCTest do
 
     test "no return" do
       assert_raise CompileError,
-                   "example.exs: Function call @Elixir.InvalidLocalCall.dummy does not return a value",
+                   ~r"example.exs: Function call @Elixir.InvalidLocalCall.dummy does not return a value",
                    fn ->
                      quote do
                        defmodule InvalidLocalCall do
diff --git a/test/vec_add_test.exs b/test/vec_add_test.exs
index 558bb9b..1c65bc5 100644
--- a/test/vec_add_test.exs
+++ b/test/vec_add_test.exs
@@ -8,13 +8,13 @@ defmodule VecAddTest do
   end
 
   test "wrong num of init values" do
-    assert_raise ArgumentError, "expected 8 values, got 6", fn ->
+    assert_raise CompileError, ~r"expected 8 values, got 6", fn ->
       defmodule SixInitValues do
         use Charms
         alias Charms.SIMD
 
         defm six(env, a, b, error) do
-          SIMD.new(i32(), 8).(1, 1, 1, 1, 1, 1)
+          SIMD.new(SIMD.t(i32(), 8), [1, 1, 1, 1, 1, 1])
         end
       end
     end