feat: Ops.throw

deps/ReactantExtra/API.cpp

@@ -137,6 +137,12 @@
 
 #include "llvm/Support/ExtensibleRTTI.h"
 
+// XLA FFI
+#include "xla/ffi/api/c_api.h"
+#include "xla/ffi/api/ffi.h"
+#include "xla/ffi/ffi_api.h"
+#include "xla/service/custom_call_target_registry.h"
+
 using namespace mlir;
 using namespace llvm;
 using namespace xla;
@@ -2281,3 +2287,39 @@ ifrt_loaded_executable_num_devices(ifrt::LoadedExecutable *exec) {
 }
 
 #pragma endregion
+
+// Register an XLA FFI for throwing runtime errors
+xla::ffi::Error xla_throw_error(xla::ffi::BufferR0<xla::ffi::PRED> cond,
+                                std::string_view message) {
+  if (cond.typed_data()[0])
+    return xla::ffi::Error(xla::ffi::ErrorCode::kInternal,
+                           std::string(message));
+  return xla::ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER(xla_throw_error_handler, xla_throw_error,
+                       xla::ffi::Ffi::Bind()
+                           .Arg<xla::ffi::BufferR0<xla::ffi::PRED>>()
+                           .Attr<std::string_view>("message"));
+
+XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "xla_throw_error", "Host",
+                         xla_throw_error_handler);
+
+#ifdef REACTANT_CUDA
+#include "third_party/gpus/cuda/include/cuda.h"
+
+xla::ffi::Error xla_throw_error(cudaStream_t stream,
+                                xla::ffi::BufferR0<xla::ffi::PRED> cond,
+                                std::string_view message) {
+  return xla_throw_error(cond, message);
+}
+
+XLA_FFI_DEFINE_HANDLER(xla_throw_error_handler_cuda, xla_throw_error,
+                       xla::ffi::Ffi::Bind()
+                           .Ctx<xla::ffi::PlatformStream<CUstream>>()
+                           .Arg<xla::ffi::BufferR0<xla::ffi::PRED>>()
+                           .Attr<std::string_view>("message"));
+
+XLA_FFI_REGISTER_HANDLER(xla::ffi::GetXlaFfiApi(), "xla_throw_error", "CUDA",
+                         xla_throw_error_handler_cuda);
+#endif

src/Ops.jl

@@ -1074,9 +1074,7 @@ end
     @assert fn_name == "comparator" "$comparator: no function generated"
     ftype_attr = MLIR.IR.attr(func, "function_type")
     ftype = MLIR.IR.Type(ftype_attr)
-    @assert MLIR.IR.result(ftype) == MLIR.IR.TensorType((), MLIR.IR.Type(Bool)) error(
-        "$comparator return type is not tensor<i1>"
-    )
+    @assert MLIR.IR.result(ftype) == MLIR.IR.TensorType((), MLIR.IR.Type(Bool)) "$comparator return type is not tensor<i1>"
 
     comparator = MLIR.IR.Region()
     MLIR.API.mlirRegionTakeBody(comparator, MLIR.IR.region(func, 1))
@@ -2313,4 +2311,33 @@ Produces a [`Reactant.MLIR.Dialects.sdy.sharding_constraint`](@ref) operation wi
     end
 end
 
+"""
+    throw(
+        msg::String,
+        condition::Union{TracedRNumber{Bool},Nothing}=nothing;
+        location=mlir_stacktrace("throw", @__FILE__, @__LINE__)
+    )
+
+Throw a runtime error with the given `msg` if `condition` is `true`. If `condition` is not provided, it defaults to `true`.
+"""
+@noinline function throw(
+    msg::String,
+    condition::Union{TracedRNumber{Bool},Nothing}=nothing;
+    location=mlir_stacktrace("throw", @__FILE__, @__LINE__)
+)
+    if condition === nothing
+        condition = Reactant.TracedUtils.promote_to(TracedRNumber{Bool}, true)
+    end
+
+    return stablehlo.custom_call(
+        MLIR.IR.Value[condition.mlir_data];
+        result_0=MLIR.IR.Type[],
+        has_side_effect=true,
+        call_target_name="xla_throw_error",
+        backend_config=MLIR.IR.Attribute(Dict("message" => MLIR.IR.Attribute(msg))),
+        api_version=MLIR.IR.Attribute(Int32(4)),
+        location,
+    )
+end
+
 end # module Ops

src/mlir/IR/Attribute.jl

@@ -96,7 +96,7 @@ isdict(attr::Attribute) = API.mlirAttributeIsADictionary(attr)
 Creates a dictionary attribute containing the given list of elements in the provided context.
 """
 function Attribute(attrs::Dict; context::Context=context())
-    attrs = map(splat(NamedAttribute), attrs)
+    attrs = [NamedAttribute(k, v) for (k, v) in attrs]
     return Attribute(API.mlirDictionaryAttrGet(context, length(attrs), attrs))
 end