Function bodies 219 total

fine::Ok<> runtime_callback_reply(ErlNifEnv *env,
                                  fine::ResourcePtr<Pending> pending,
                                  fine::Atom status, fine::Term result) {
  deliver_reply(env, pending, status, result);
  return fine::Ok();
}

void deliver_reply(ErlNifEnv *env, fine::ResourcePtr<Pending> pending,
                   fine::Atom status, fine::Term result_term) {
  Result cb_result;

  if (status == "ok") {
    // Successful reply: result_term is a list of binaries that we decode into
    // raw byte vectors via Fine and copy directly into the registered output
    // buffers.
    try {
      auto payloads = fine::decode<std::vector<ErlNifBinary>>(env, result_term);

      std::lock_guard<std::mutex> lock(pending->mu);

      if (payloads.size() != pending->outputs.size()) {
        cb_result.ok = false;
        cb_result.error =
            "mismatched number of callback outputs vs registered buffers";
      } else {
        cb_result.ok = true;

        for (size_t i = 0; i < payloads.size(); ++i) {
          const ErlNifBinary &bytes = payloads[i];
          auto &out_buf = pending->outputs[i];

          if (bytes.size != out_buf.size) {
            cb_result.ok = false;
            cb_result.error =
      

Result InvokeRuntimeCallback(
    xla::ffi::Span<const int64_t> callback_id_words, uint64_t callback_id_size,
    const Arg &callback_server_pid_arg, const std::vector<Arg> &inputs,
    const std::vector<OutputBuffer> &outputs) {
  const int64_t callback_server_pid_dim = static_cast<int64_t>(sizeof(ErlNifPid));
  const size_t callback_server_pid_size = sizeof(ErlNifPid);

  auto pending = fine::make_resource<Pending>(outputs);

  ErlNifEnv *msg_env = enif_alloc_env();

  // Reinterpret the 64-bit words as a contiguous byte buffer and use the
  // original (unpadded) size when decoding the callback id term.
  if (callback_id_size > callback_id_words.size() * sizeof(int64_t)) {
    Result res;
    res.ok = false;
    res.error = "inconsistent callback id size";
    return res;
  }

  if (callback_server_pid_arg.dtype != xla::ffi::DataType::U8) {
    Result res;
    res.ok = false;
    res.error = "callback server pid tensor must have dtype u8";
    return res;
  }

  if (callback_server

  static ERL_NIF_TERM encode(ErlNifEnv *env, const xla::ffi::DataType &dtype) {
    using DT = xla::ffi::DataType;

    // Tokens are encoded as the atom :token with empty shape.
    if (dtype == DT::TOKEN) {
      return fine::encode(env, exla::atoms::token);
    }

    std::optional<fine::Atom> type_name;
    std::optional<uint64_t> type_size;

    switch (dtype) {
    case DT::PRED:
      type_name = exla::atoms::pred;
      type_size = 8;
      break;

    case DT::U8:
      type_name = exla::atoms::u;
      type_size = 8;
      break;
    case DT::U16:
      type_name = exla::atoms::u;
      type_size = 16;
      break;
    case DT::U32:
      type_name = exla::atoms::u;
      type_size = 32;
      break;
    case DT::U64:
      type_name = exla::atoms::u;
      type_size = 64;
      break;

    case DT::S8:
      type_name = exla::atoms::s;
      type_size = 8;
      break;
    case DT::S16:
      type_name = exla::atoms::s;
      type_size = 16;
      break;
    case DT::S32:
  

ffi::Error
exla_runtime_callback_impl(ffi::RemainingArgs args,
                           ffi::Span<const int64_t> callback_id_words,
                           uint64_t callback_id_size, ffi::RemainingRets rets) {
  if (args.size() == 0) {
    return ffi::Error(ffi::ErrorCode::kInternal,
                      "runtime callback missing callback server pid operand");
  }

  const size_t callback_args_end = args.size();

  // Collect all input tensors into lightweight payload views.
  std::vector<exla::callback_bridge::Arg> inputs;
  inputs.reserve(callback_args_end - 1);
  exla::callback_bridge::Arg callback_server_pid_arg;

  for (size_t i = 0; i < args.size(); ++i) {
    auto maybe_buf_or = args.get<ffi::AnyBuffer>(i);
    if (!maybe_buf_or) {
      return maybe_buf_or.error();
    }

    ffi::AnyBuffer buf = *maybe_buf_or;

    exla::callback_bridge::Arg tensor;
    tensor.dtype = buf.element_type();

    auto dims = buf.dimensions();
    tensor.dims.assign(dims.begin(), dims.end())

ffi::Error exla_runtime_callback_cuda_impl(
    CUstream stream, ffi::RemainingArgs args,
    ffi::Span<const int64_t> callback_id_words, uint64_t callback_id_size,
    ffi::RemainingRets rets) {
  if (args.size() == 0) {
    return ffi::Error(ffi::ErrorCode::kInternal,
                      "runtime callback missing callback server pid operand");
  }

  const size_t callback_args_end = args.size();

  // Keep host buffers alive for the duration of the callback.
  std::vector<std::vector<uint8_t>> host_input_buffers;
  host_input_buffers.reserve(callback_args_end);

  std::vector<exla::callback_bridge::Arg> inputs;
  inputs.reserve(callback_args_end - 1);
  exla::callback_bridge::Arg callback_server_pid_arg;

  for (size_t i = 0; i < args.size(); ++i) {
    auto maybe_buf_or = args.get<ffi::AnyBuffer>(i);
    if (!maybe_buf_or) {
      return maybe_buf_or.error();
    }

    ffi::AnyBuffer buf = *maybe_buf_or;
    size_t size_bytes = buf.size_bytes();

    host_input_buffers.emplace_b

ffi::Error exla_runtime_callback_cuda_stub(
    ffi::RemainingArgs, ffi::Span<const int64_t>, uint64_t,
    ffi::RemainingRets) {
  return ffi::Error(ffi::ErrorCode::kUnimplemented,
                    "EXLA was not compiled with CUDA support. This error means your EXLA compilation is out of sync with your libexla.so NIF.");
}

fine::ResourcePtr<ExlaBuffer> decode_exla_buffer(ErlNifEnv *env,
                                                 fine::Term buffer_term) {
  try {
    return fine::decode<fine::ResourcePtr<ExlaBuffer>>(env, buffer_term);
  } catch (std::invalid_argument) {
    throw std::invalid_argument(
        "unable to get buffer. It may belong to another node, "
        "consider using Nx.backend_transfer/1");
  }
}

fine::ResourcePtr<llvm::StdThreadPool>
mlir_new_thread_pool(ErlNifEnv *env, int64_t concurrency) {
  auto strategy = llvm::hardware_concurrency(concurrency);
  return fine::make_resource<llvm::StdThreadPool>(strategy);
}

fine::ResourcePtr<mlir::MLIRContext>
mlir_new_context(ErlNifEnv *env,
                 fine::ResourcePtr<llvm::StdThreadPool> thread_pool) {
  auto context = fine::make_resource<mlir::MLIRContext>(
      mlir::MLIRContext::Threading::DISABLED);

  context->setThreadPool(*thread_pool);
  context->getOrLoadDialect<mlir::func::FuncDialect>();
  context->getOrLoadDialect<mlir::stablehlo::StablehloDialect>();
  context->getOrLoadDialect<mlir::chlo::ChloDialect>();
  context->getOrLoadDialect<mlir::sdy::SdyDialect>();
  mlir::sdy::registerAllDialects(
      const_cast<mlir::DialectRegistry &>(context->getDialectRegistry()));

  return context;
}

fine::ResourcePtr<MLIRModule>
mlir_new_module(ErlNifEnv *env, fine::ResourcePtr<mlir::MLIRContext> ctx) {
  return fine::make_resource<MLIRModule>(ctx);
}

fine::ResourcePtr<MLIRFunction> mlir_create_function(
    ErlNifEnv *env, fine::ResourcePtr<MLIRModule> module, std::string func_name,
    std::vector<std::string> arg_type_strings,
    std::vector<std::string> ret_type_strings, bool is_public) {
  auto arg_types = std::vector<mlir::Type>{};

  for (auto const &type_string : arg_type_strings) {
    auto type = module->ParseType(type_string);
    arg_types.push_back(type);
  }

  auto ret_types = std::vector<mlir::Type>{};

  for (auto const &type_string : ret_type_strings) {
    auto type = module->ParseType(type_string);
    ret_types.push_back(type);
  }

  auto func_op =
      module->CreateFunction(func_name, arg_types, ret_types, is_public);
  return fine::make_resource<MLIRFunction>(module, std::move(func_op));
}

std::vector<fine::ResourcePtr<mlir::Value>>
mlir_get_function_arguments(ErlNifEnv *env,
                            fine::ResourcePtr<MLIRFunction> function) {
  auto args = function->GetArguments();
  std::vector<fine::ResourcePtr<mlir::Value>> values;
  values.reserve(args.size());

  for (const auto &arg : args) {
    values.push_back(fine::make_resource<mlir::Value>(arg));
  }

  return values;
}

std::vector<fine::ResourcePtr<mlir::Value>>
mlir_op(ErlNifEnv *env, fine::ResourcePtr<MLIRFunction> function,
        std::string op_name,
        std::vector<fine::ResourcePtr<mlir::Value>> operands,
        std::vector<std::string> result_type_strings,
        std::vector<std::tuple<fine::Atom, std::string>> attributes_kwlist,
        std::vector<fine::ResourcePtr<mlir::Region>> regions) {
  auto result_types = std::vector<mlir::Type>{};

  for (auto const &type_string : result_type_strings) {
    auto type = function->module()->ParseType(type_string);
    result_types.push_back(type);
  }

  auto attributes = std::vector<std::tuple<std::string, mlir::Attribute>>{};

  for (auto const &[key, value] : attributes_kwlist) {
    auto attribute_value = function->module()->ParseAttribute(value);
    attributes.push_back(std::make_tuple(key.to_string(), attribute_value));
  }

  return function->Op(op_name, operands, result_types, attributes, regions);
}

std::tuple<fine::ResourcePtr<mlir::Region>,
           std::vector<fine::ResourcePtr<mlir::Value>>>
mlir_push_region(ErlNifEnv *env, fine::ResourcePtr<MLIRFunction> function,
                 std::vector<std::string> arg_types) {
  auto types = std::vector<mlir::Type>{};

  for (auto const &type_string : arg_types) {
    auto type = function->module()->ParseType(type_string);
    types.push_back(type);
  }

  return function->PushRegion(types);
}

fine::Ok<> mlir_pop_region(ErlNifEnv *env,
                           fine::ResourcePtr<MLIRFunction> function) {
  function->PopRegion();
  return fine::Ok();
}

fine::Ok<> mlir_add_mesh(ErlNifEnv *env, fine::ResourcePtr<MLIRModule> module,
                         std::string mesh_name,
                         std::vector<std::tuple<std::string, int64_t>> axes) {
  auto builder = module->builder();
  auto context = module->module()->getContext();

  llvm::SmallVector<mlir::sdy::MeshAxisAttr> axis_attrs;
  for (auto [name, size] : axes) {
    axis_attrs.push_back(mlir::sdy::MeshAxisAttr::get(context, name, size));
  }

  auto mesh_attr = mlir::sdy::MeshAttr::get(context, axis_attrs);

  // Create the mesh op at the beginning of the module
  auto module_op = module->module();
  auto &body_region = module_op.getBodyRegion();
  mlir::OpBuilder::InsertionGuard guard(*builder);
  builder->setInsertionPointToStart(&body_region.front());

  mlir::OperationState state(builder->getUnknownLoc(), "sdy.mesh");
  mlir::sdy::MeshOp::build(*builder, state, mesh_name, mesh_attr);
  builder->create(state);

  return fine::Ok();
}

mlir::sdy::TensorShardingAttr mlir_create_tensor_sharding_attr(
    mlir::MLIRContext *context, std::string mesh_name,
    std::vector<std::vector<std::string>> dim_shardings) {
  llvm::SmallVector<mlir::sdy::DimensionShardingAttr> dim_sharding_attrs;
  for (const auto &dim : dim_shardings) {
    llvm::SmallVector<mlir::sdy::AxisRefAttr> axis_refs;
    for (const auto &axis : dim) {
      axis_refs.push_back(mlir::sdy::AxisRefAttr::get(context, axis));
    }
    dim_sharding_attrs.push_back(mlir::sdy::DimensionShardingAttr::get(
        context, axis_refs, /*is_closed=*/false, /*priority=*/0));
  }

  return mlir::sdy::TensorShardingAttr::get(
      context, mesh_name, dim_sharding_attrs,
      /*replicated_axes=*/llvm::ArrayRef<mlir::sdy::AxisRefAttr>(),
      /*unreduced_axes=*/llvm::ArrayRef<mlir::sdy::AxisRefAttr>());
}

fine::Ok<> mlir_set_function_argument_attribute(
    ErlNifEnv *env, fine::ResourcePtr<MLIRFunction> function, int64_t arg_index,
    std::string attribute_name, std::string mesh_name,
    std::vector<std::vector<std::string>> dim_shardings) {

  auto context = function->module()->module()->getContext();
  auto sharding_attr =
      mlir_create_tensor_sharding_attr(context, mesh_name, dim_shardings);

  function->function().setArgAttr(arg_index, attribute_name, sharding_attr);

  return fine::Ok();
}

mlir::Type mlir_get_typespec(ErlNifEnv *env,
                             fine::ResourcePtr<mlir::Value> value) {
  return value->getType();
}

std::string mlir_module_to_string(ErlNifEnv *env,
                                  fine::ResourcePtr<MLIRModule> module) {
  return module->ToString();
}

template <typename T> T unwrap(tsl::StatusOr<T> status_or) {
  if (!status_or.ok()) {
    throw std::runtime_error(status_or.status().message().data());
  }

  return std::move(status_or.value());
}

void unwrap(tsl::Status status) {
  if (!status.ok()) {
    throw std::runtime_error(status.message().data());
  }
}

fine::ResourcePtr<ExlaExecutable>
mlir_compile(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
             fine::ResourcePtr<MLIRModule> module,
             std::vector<xla::Shape> argument_layouts, int64_t num_replicas,
             int64_t num_partitions, bool use_spmd, int64_t device_id) {
  auto build_options = xla::ExecutableBuildOptions();

  build_options.set_num_replicas(num_replicas);
  build_options.set_num_partitions(num_partitions);
  build_options.set_use_spmd_partitioning(use_spmd);

  // Enable Shardy partitioner when using SPMD to support output sharding
  // Shardy runs propagation before the SPMD partitioner and properly handles
  // FuncResultSharding custom calls that would otherwise fail the side-effect
  // check
  if (use_spmd) {
    build_options.set_use_shardy_partitioner(true);
  }

  auto compile_portable_executable = false;
  if (device_id >= 0) {
    compile_portable_executable = true;
    build_options.set_device_ordinal(device_id);
  }

  return 

std::variant<std::tuple<fine::Atom, uint64_t, uint64_t>,
             std::tuple<fine::Atom, std::string, uint64_t>>
get_buffer_device_pointer(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
                          fine::Term buffer_term, fine::Atom pointer_kind,
                          int64_t shm_permissions) {
  auto buffer = decode_exla_buffer(env, buffer_term);

  uint64_t device_size = unwrap(buffer->GetOnDeviceSizeInBytes());
  uint64_t ptr = unwrap(buffer->GetDevicePointer(client->client()));

  if (pointer_kind == "local") {
    return std::make_tuple(pointer_kind, ptr, device_size);
  }

  if (pointer_kind == "host_ipc") {
    auto handle_name =
        "exla:ipc:" + std::to_string(getpid()) + ":" + std::to_string(ptr);
    auto fd = get_ipc_handle(handle_name.c_str(), device_size,
                             static_cast<mode_t>(shm_permissions));

    if (fd == -1) {
      throw std::runtime_error("unable to get IPC handle");
    }

    auto ipc_ptr = open_ipc_ha

fine::ResourcePtr<ExlaBuffer> create_buffer_from_device_pointer(
    ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
    fine::Atom pointer_kind, fine::Term pointer_data, xla::Shape shape,
    int64_t device_id) {
  void *ptr = nullptr;
  std::function<void()> on_delete_callback = []() {};

  if (pointer_kind == "cuda_ipc") {
    auto cuda_ipc_handle_bin = fine::decode<ErlNifBinary>(env, pointer_data);
    auto maybe_pointer = get_pointer_for_ipc_handle(
        cuda_ipc_handle_bin.data, cuda_ipc_handle_bin.size, device_id);
    if (!maybe_pointer) {
      throw std::runtime_error("unable to get pointer for IPC handle");
    }
    ptr = maybe_pointer.value();
  } else if (pointer_kind == "host_ipc") {
    auto memname = fine::decode<std::string>(env, pointer_data);
    auto device_size = xla::ShapeUtil::ByteSizeOf(shape);
    int writable = 0;
    auto fd = open_existing_ipc_handle(memname.c_str(), &writable);
    if (fd == -1) {
      throw std::runtime_error("unable to get fd 

fine::ResourcePtr<ExlaBuffer>
binary_to_device_mem(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
                     fine::Term data, xla::Shape shape, int64_t device_id) {
  return unwrap(client->BufferFromBinary(data, shape, device_id));
}

fine::Term read_device_mem(ErlNifEnv *env, fine::Term buffer_term,
                           int64_t size) {
  auto buffer = decode_exla_buffer(env, buffer_term);
  return unwrap(buffer->ToBinary(env, size));
}

xla::Shape get_buffer_typespec(ErlNifEnv *env, fine::Term buffer_term) {
  auto buffer = decode_exla_buffer(env, buffer_term);
  return unwrap(buffer->buffer()->logical_on_device_shape());
}

std::variant<fine::Ok<>, fine::Error<fine::Atom>>
deallocate_device_mem(ErlNifEnv *env, fine::Term buffer_term) {
  auto buffer = decode_exla_buffer(env, buffer_term);

  tsl::Status dealloc_status = buffer->Deallocate();

  if (!dealloc_status.ok()) {
    return fine::Error(atoms::already_deallocated);
  } else {
    return fine::Ok();
  }
}

fine::Ok<> transfer_to_infeed(ErlNifEnv *env,
                              fine::ResourcePtr<ExlaClient> client,
                              int64_t device_id,
                              std::vector<ErlNifBinary> buffers,
                              std::vector<xla::Shape> shapes) {
  unwrap(client->TransferToInfeed(env, buffers, shapes, device_id));

  return fine::Ok();
}

fine::Ok<> transfer_from_outfeed(ErlNifEnv *env,
                                 fine::ResourcePtr<ExlaClient> client,
                                 int64_t device_id,
                                 std::vector<xla::Shape> shapes, ErlNifPid pid,
                                 fine::Term ref) {
  for (auto &shape : shapes) {
    auto msg_env = enif_alloc_env();
    auto msg = unwrap(client->TransferFromOutfeed(msg_env, device_id, shape));
    enif_send(env, &pid, msg_env, enif_make_tuple(msg_env, 2, ref, msg));
    enif_free_env(msg_env);
  }

  return fine::Ok();
}

fine::ResourcePtr<ExlaBuffer>
copy_buffer_to_device(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
                      fine::Term buffer_term, int64_t device_id) {
  auto buffer = decode_exla_buffer(env, buffer_term);

  auto device = unwrap(
      client->client()->LookupDevice(xla::PjRtGlobalDeviceId(device_id)));

  return unwrap(buffer->CopyToDevice(device));
}

fine::ResourcePtr<ExlaClient> get_host_client(ErlNifEnv *env) {
  return unwrap(GetHostClient());
}

fine::ResourcePtr<ExlaClient>
get_gpu_client(ErlNifEnv *env, double memory_fraction, bool preallocate) {
  return unwrap(GetGpuClient(memory_fraction, preallocate,
                             xla::GpuAllocatorConfig::Kind::kBFC));
}

fine::ResourcePtr<ExlaClient> get_tpu_client(ErlNifEnv *env) {
  return unwrap(GetTpuClient());
}

fine::ResourcePtr<ExlaClient> get_c_api_client(ErlNifEnv *env,
                                               std::string device_type) {
  return unwrap(GetCApiClient(device_type));
}

fine::Ok<> load_pjrt_plugin(ErlNifEnv *env, std::string device_type,
                            std::string library_path) {
  unwrap(pjrt::LoadPjrtPlugin(device_type, library_path));
  return fine::Ok();
}

int64_t get_device_count(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client) {
  return client->client()->device_count();
}

std::map<fine::Atom, int64_t> get_supported_platforms(ErlNifEnv *env) {
  auto platforms = unwrap(xla::PlatformUtil::GetSupportedPlatforms());

  std::map<fine::Atom, int64_t> platform_info;

  for (auto &platform : platforms) {
    auto key = fine::Atom(absl::AsciiStrToLower(platform->Name()));
    auto device_count = platform->VisibleDeviceCount();
    platform_info.insert({key, device_count});
  }

  return platform_info;
}

ExlaExecutable::RunResult run(ErlNifEnv *env,
                              fine::ResourcePtr<ExlaExecutable> executable,
                              ExlaExecutable::RunArguments arguments,
                              int64_t device_id,
                              std::optional<ErlNifPid> callback_server_pid) {
  auto result = unwrap(executable->Run(env, arguments, device_id));

  if (callback_server_pid.has_value()) {
    auto pid = callback_server_pid.value();
    auto msg_env = enif_alloc_env();
    enif_send(msg_env, &pid, msg_env, fine::encode(msg_env, exla::atoms::stop));
    enif_free_env(msg_env);
  }

  return result;
}

ExlaExecutable::RunResult
run_cpu(ErlNifEnv *env, fine::ResourcePtr<ExlaExecutable> executable,
        ExlaExecutable::RunArguments arguments, int64_t device_id,
        std::optional<ErlNifPid> callback_server_pid) {
  return run(env, executable, arguments, device_id, callback_server_pid);
}

ExlaExecutable::RunResult run_io(ErlNifEnv *env,
                                 fine::ResourcePtr<ExlaExecutable> executable,
                                 ExlaExecutable::RunArguments arguments,
                                 int64_t device_id,
                                 std::optional<ErlNifPid> callback_server_pid) {
  return run(env, executable, arguments, device_id, callback_server_pid);
}

std::string serialize_executable(ErlNifEnv *env,
                                 fine::ResourcePtr<ExlaExecutable> executable) {
  return unwrap(executable->SerializeExecutable());
}

fine::ResourcePtr<ExlaExecutable>
deserialize_executable(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client,
                       std::string serialized) {
  return unwrap(client->DeserializeExecutable(serialized));
}

int64_t get_allocated_memory(ErlNifEnv *env,
                             fine::ResourcePtr<ExlaClient> client) {
  return client->GetAllocatedMemory();
}

int64_t get_peak_memory(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client) {
  return client->GetPeakMemory();
}

fine::Ok<> reset_peak_memory(ErlNifEnv *env,
                             fine::ResourcePtr<ExlaClient> client) {
  client->ResetPeakMemory();
  return fine::Ok();
}

std::map<int64_t, int64_t>
get_per_device_memory(ErlNifEnv *env, fine::ResourcePtr<ExlaClient> client) {
  auto device_memory = client->GetPerDeviceMemory();
  std::map<int64_t, int64_t> result;
  for (const auto &pair : device_memory) {
    result[pair.first] = pair.second;
  }
  return result;
}

int64_t callback_server_pid_size(ErlNifEnv *env) {
  (void)env;
  return static_cast<int64_t>(sizeof(ErlNifPid));
}