Add runtime raise_if to Nx.Defn.Kernel

exla/lib/exla/defn.ex

@@ -246,28 +246,56 @@ defmodule EXLA.Defn do
       end)
 
     infeeds = Map.new(infeeds)
+    error_ref = make_ref()
 
     {:ok, outfeed_pid} =
-      Outfeed.start_child(executable, outfeed, Process.group_leader(), infeeds)
+      Outfeed.start_child(
+        executable,
+        outfeed,
+        Process.group_leader(),
+        infeeds,
+        {self(), error_ref}
+      )
 
     ref = Process.monitor(outfeed_pid)
 
     run_options = Keyword.put(run_options, :callback_server_pid, outfeed_pid)
 
     {:ok, runner} =
       EXLA.Defn.Runner.start_link(lock, fn ->
-        EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options)
+        try do
+          EXLA.Executable.run(executable, [Enum.reverse(buffers)], run_options)
+        after
+          send(outfeed_pid, :stop)
+        end
       end)
 
     _ = EXLA.Defn.Lock.transfer(lock, fn -> send(runner, lock) end, outfeed_pid)
 
     receive do
       {:DOWN, ^ref, _, _, _} ->
-        results = EXLA.Defn.Runner.read(runner)
+        runner_result = EXLA.Defn.Runner.read(runner)
+
+        callback_error =
+          receive do
+            {:exla_callback_error, ^error_ref, kind, reason, stacktrace} ->
+              {kind, reason, stacktrace}
+          after
+            0 -> nil
+          end
 
-        Enum.map(results, fn result ->
-          EXLA.Defn.Buffers.to_nx!(result, outputs, executable.mesh)
-        end)
+        case {callback_error, runner_result} do
+          {{kind, reason, stacktrace}, _runner_result} ->
+            :erlang.raise(kind, reason, stacktrace)
+
+          {nil, {:error, kind, reason, stacktrace}} ->
+            :erlang.raise(kind, reason, stacktrace)
+
+          {nil, {:ok, results}} ->
+            Enum.map(results, fn result ->
+              EXLA.Defn.Buffers.to_nx!(result, outputs, executable.mesh)
+            end)
+        end
     end
   end
 

exla/lib/exla/defn/outfeed.ex

@@ -216,7 +216,8 @@ defmodule EXLA.Defn.Outfeed do
         %EXLA.Executable{} = executable,
         %Outfeed{} = outfeed,
         group_leader,
-        infeeds
+        infeeds,
+        error_sink
       ) do
     %{client: client, device_id: device_id} = executable
 
@@ -232,22 +233,32 @@ defmodule EXLA.Defn.Outfeed do
     callbacks = resolve_callbacks(callbacks, io_calls)
 
     Task.Supervisor.start_child(EXLA.Defn.TaskSupervisor, fn ->
-      init(client, device_id, infeed_flags, infeeds, callbacks, group_leader)
+      init(client, device_id, infeed_flags, infeeds, callbacks, group_leader, error_sink)
     end)
   end
 
-  defp init(client, device_id, infeed_flags, infeeds, callbacks, group_leader) do
+  defp init(client, device_id, infeed_flags, infeeds, callbacks, group_leader, error_sink) do
     Process.flag(:trap_exit, true)
     # Copy the group leader so we report to the proper device
     Process.group_leader(self(), group_leader)
 
     ref = make_ref()
     typespec = EXLA.Typespec.tensor({:u, 16}, {})
 
-    loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks)
+    loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks, error_sink, nil)
   end
 
-  defp loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks) do
+  defp loop(
+         client,
+         device_id,
+         ref,
+         typespec,
+         infeed_flags,
+         infeeds,
+         callbacks,
+         error_sink,
+         callback_error
+       ) do
     if active_infeed_flags?(infeed_flags) do
       :ok = EXLA.Client.from_outfeed(client, device_id, [typespec], self(), ref)
     end
@@ -262,7 +273,9 @@ defmodule EXLA.Defn.Outfeed do
           typespec,
           drop_infeed_flags(infeed_flags),
           infeeds,
-          callbacks
+          callbacks,
+          error_sink,
+          callback_error
         )
 
       {^ref, <<flag::native-unsigned-16>>} ->
@@ -276,22 +289,61 @@ defmodule EXLA.Defn.Outfeed do
 
             EXLA.Client.to_infeed(client, device_id, [{data, data_typespec}])
 
-            loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks)
+            loop(
+              client,
+              device_id,
+              ref,
+              typespec,
+              infeed_flags,
+              infeeds,
+              callbacks,
+              error_sink,
+              callback_error
+            )
         end
 
       {:exla_runtime_call, callback_id, args_spec, reply_tag} ->
-        send_callback_reply(callbacks, callback_id, args_spec, reply_tag)
-        loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks)
+        error = send_callback_reply(callbacks, callback_id, args_spec, reply_tag)
+
+        loop(
+          client,
+          device_id,
+          ref,
+          typespec,
+          infeed_flags,
+          infeeds,
+          callbacks,
+          error_sink,
+          callback_error || error
+        )
 
       :stop ->
+        maybe_send_callback_error(error_sink, callback_error)
         :ok
 
       other ->
         Logger.debug("EXLA.Outfeed ignoring unexpected message: #{inspect(other)}")
-        loop(client, device_id, ref, typespec, infeed_flags, infeeds, callbacks)
+
+        loop(
+          client,
+          device_id,
+          ref,
+          typespec,
+          infeed_flags,
+          infeeds,
+          callbacks,
+          error_sink,
+          callback_error
+        )
     end
   end
 
+  defp maybe_send_callback_error({pid, ref}, {kind, reason, stacktrace}) do
+    send(pid, {:exla_callback_error, ref, kind, reason, stacktrace})
+  end
+
+  defp maybe_send_callback_error(_error_sink, nil), do: :ok
+
   defp drop_infeed_flags(infeed_flags) do
     keys = for {k, _} <- infeed_flags, is_integer(k), do: k
     Map.drop(infeed_flags, keys)
@@ -315,35 +367,45 @@ defmodule EXLA.Defn.Outfeed do
   end
 
   defp send_callback_reply(callbacks, callback_id, args_spec, reply_tag) do
-    reply =
+    {reply, callback_error} =
       try do
-        with {:ok, callback} <- Map.fetch(callbacks, callback_id),
-             {:ok, tensor_args} <- materialize_callback_args(callback, args_spec) do
-          callback
-          |> invoke_callback(tensor_args)
-          |> encode_callback_reply()
-        else
-          :error ->
-            Logger.error(
-              "EXLA.Outfeed received callback id #{inspect(callback_id)} that is not registered"
-            )
+        result =
+          with {:ok, callback} <- Map.fetch(callbacks, callback_id),
+               {:ok, tensor_args} <- materialize_callback_args(callback, args_spec) do
+            invoke_callback(callback, tensor_args)
+          else
+            :error ->
+              Logger.error(
+                "EXLA.Outfeed received callback id #{inspect(callback_id)} that is not registered"
+              )
 
-            encode_callback_reply({:error, :unknown_callback})
+              {:error, :unknown_callback}
 
-          {:error, _} = error ->
-            error
-        end
+            {:error, _} = error ->
+              error
+          end
+
+        {encode_callback_reply(result), nil}
       rescue
         exception ->
-          send(self(), :stop)
-          {:error, {:exception, Exception.format(:error, exception, __STACKTRACE__)}}
+          stacktrace = __STACKTRACE__
+
+          {
+            {:error, {:exception, Exception.format(:error, exception, stacktrace)}},
+            {:error, exception, stacktrace}
+          }
       catch
         kind, reason ->
-          send(self(), :stop)
-          {:error, {kind, Exception.format(kind, reason, __STACKTRACE__)}}
+          stacktrace = __STACKTRACE__
+
+          {
+            {:error, {kind, Exception.format(kind, reason, stacktrace)}},
+            {kind, reason, stacktrace}
+          }
       end
 
     deliver_native_reply(reply_tag, reply)
+    callback_error
   end
 
   defp deliver_native_reply(reply_tag, reply) do
@@ -386,10 +448,10 @@ defmodule EXLA.Defn.Outfeed do
       {:ok, []}
     rescue
       exception ->
-        {:error, {:exception, Exception.format(:error, exception, __STACKTRACE__)}}
+        {:error, {:callback_error, :error, exception, __STACKTRACE__}}
     catch
       kind, reason ->
-        {:error, {kind, Exception.format(kind, reason, __STACKTRACE__)}}
+        {:error, {:callback_error, kind, reason, __STACKTRACE__}}
     end
   end
 
@@ -403,10 +465,10 @@ defmodule EXLA.Defn.Outfeed do
         end
       rescue
         exception ->
-          {:error, {:exception, exception, __STACKTRACE__}}
+          {:error, {:callback_error, :error, exception, __STACKTRACE__}}
       catch
         kind, reason ->
-          {:error, {kind, reason}}
+          {:error, {:callback_error, kind, reason, __STACKTRACE__}}
       end
 
     case result do
@@ -468,8 +530,8 @@ defmodule EXLA.Defn.Outfeed do
     raise RuntimeError.exception(msg)
   end
 
-  defp encode_callback_reply({:error, {:exception, exception, _stack}}) do
-    raise exception
+  defp encode_callback_reply({:error, {:callback_error, kind, reason, stacktrace}}) do
+    :erlang.raise(kind, reason, stacktrace)
   end
 
   defp encode_callback_reply({:error, {kind, reason}}) when is_binary(reason) do

exla/lib/exla/defn/runner.ex

@@ -21,7 +21,16 @@ defmodule EXLA.Defn.Runner do
   @impl true
   def handle_continue({ref, fun}, nil) do
     receive do
-      ^ref -> {:noreply, fun.()}
+      ^ref ->
+        result =
+          try do
+            {:ok, fun.()}
+          catch
+            kind, reason ->
+              {:error, kind, reason, __STACKTRACE__}
+          end
+
+        {:noreply, result}
     end
   end
 

exla/test/exla/defn/api_test.exs

@@ -294,15 +294,10 @@ defmodule EXLA.Defn.APITest do
       io_call(a + b, :raises, fn _ -> raise "boom" end)
     end
 
-    @tag :capture_log
-    test "halts outfeed when io_call raises" do
-      {_pid, ref} =
-        spawn_monitor(fn ->
-          EXLA.jit(&hook_raises/2).(2, 3)
-        end)
-
-      assert_receive {:DOWN, ^ref, :process, _, {%RuntimeError{message: message}, _}}
-      assert message =~ "boom"
+    test "propagates the original exception when io_call raises" do
+      assert_raise RuntimeError, "boom", fn ->
+        EXLA.jit(&hook_raises/2).(2, 3)
+      end
     end
 
     defn side_effect_hooks(a, b) do
@@ -396,6 +391,62 @@ defmodule EXLA.Defn.APITest do
     end
   end
 
+  describe "raise_if" do
+    defmodule RaiseIfError do
+      defexception [:message, :value]
+    end
+
+    defn runtime_raise_if(value, predicate) do
+      raise_if(value, predicate, "runtime check failed")
+    end
+
+    defn custom_runtime_raise_if(value, predicate) do
+      raise_if(value, predicate, RaiseIfError,
+        message: "custom runtime check failed",
+        value: :preserved
+      )
+    end
+
+    defn halt_on_nth(x, n) do
+      {x, i, _n} =
+        while {x, i = 0, n}, i < 10 do
+          i = raise_if(i, i == n, "Halting on selected iteration")
+          {x + 1, i + 1, n}
+        end
+
+      {x, i}
+    end
+
+    test "passes values through when the predicate is false" do
+      assert_equal(EXLA.jit(&runtime_raise_if/2).(Nx.tensor([1, 2]), 0), Nx.tensor([1, 2]))
+    end
+
+    test "raises when the predicate is true" do
+      assert_raise RuntimeError, "runtime check failed", fn ->
+        EXLA.jit(&runtime_raise_if/2).(1, 1)
+      end
+    end
+
+    test "raises custom exceptions with arguments" do
+      error =
+        assert_raise RaiseIfError, "custom runtime check failed", fn ->
+          EXLA.jit(&custom_runtime_raise_if/2).(1, 1)
+        end
+
+      assert error.value == :preserved
+    end
+
+    test "halts a while loop on the selected iteration" do
+      assert {x, i} = EXLA.jit(&halt_on_nth/2).(0, 11)
+      assert_equal(x, Nx.tensor(10))
+      assert_equal(i, Nx.tensor(10))
+
+      assert_raise RuntimeError, "Halting on selected iteration", fn ->
+        EXLA.jit(&halt_on_nth/2).(0, 5)
+      end
+    end
+  end
+
   describe "cross-client io_calls" do
     defn hooked_add(a, b) do
       io_call(a + b, :add)