Function bodies 219 total

mlir::Attribute MLIRModule::ParseAttribute(std::string string) {
  auto attribute = mlir::parseAttribute(string, context_.get());

  if (attribute == nullptr) {
    throw std::runtime_error("unable to parse MLIR type: " + string);
  }

  return attribute;
}

class MLIRFunction {
 public:
  MLIRFunction(fine::ResourcePtr<MLIRModule> module, std::unique_ptr<mlir::func::FuncOp> func);

  std::vector<fine::ResourcePtr<mlir::Value>> Op(
      std::string op_name, std::vector<fine::ResourcePtr<mlir::Value>> operands,
      std::vector<mlir::Type> result_types,
      std::vector<std::tuple<std::string, mlir::Attribute>> attributes,
      std::vector<fine::ResourcePtr<mlir::Region>> regions);

  std::tuple<fine::ResourcePtr<mlir::Region>, std::vector<fine::ResourcePtr<mlir::Value>>> PushRegion(std::vector<mlir::Type> types);
  void PopRegion();

  llvm::MutableArrayRef<mlir::BlockArgument> GetArguments() { return func_->getBody().front().getArguments(); }

  mlir::func::FuncOp function() { return *func_; }

  fine::ResourcePtr<MLIRModule> module() { return module_; }

 private:
  fine::ResourcePtr<MLIRModule> module_;
  std::unique_ptr<mlir::func::FuncOp> func_;

  std::stack<fine::ResourcePtr<mlir::Region>> region_stack;

  void setInsertionPoin

class MLIRModule {
 public:
  MLIRModule(fine::ResourcePtr<mlir::MLIRContext> context);

  std::unique_ptr<mlir::func::FuncOp> CreateFunction(
      std::string name,
      std::vector<mlir::Type> arg_types,
      std::vector<mlir::Type> ret_types,
      bool is_public);

  std::string ToString();

  // Note: ParseAttribute and ParseType return nullptr if the parsing fails
  mlir::Type ParseType(std::string);
  mlir::Attribute ParseAttribute(std::string);

  mlir::ModuleOp module() { return module_.get(); }
  mlir::OpBuilder *builder() { return builder_.get(); }

 private:
  fine::ResourcePtr<mlir::MLIRContext> context_;
  mlir::OwningOpRef<mlir::ModuleOp> module_;
  std::unique_ptr<mlir::OpBuilder> builder_;
};

  static xla::Shape decode(ErlNifEnv *env, const ERL_NIF_TERM &term) {
    ERL_NIF_TERM type_term;
    ERL_NIF_TERM shape_term;

    if (!enif_get_map_value(env, term, fine::encode(env, exla::atoms::type),
                            &type_term)) {
      throw std::invalid_argument(
          "decode failed, expected EXLA.Typespec struct");
    }

    if (!enif_get_map_value(env, term, fine::encode(env, exla::atoms::shape),
                            &shape_term)) {
      throw std::invalid_argument(
          "decode failed, expected EXLA.Typespec struct");
    }

    return xla::ShapeUtil::MakeShape(decode_type(env, type_term),
                                     decode_shape(env, shape_term));
  }

private:
  static std::vector<int64_t> decode_shape(ErlNifEnv *env,
                                           const ERL_NIF_TERM &term) {
    int size;
    const ERL_NIF_TERM *terms;

    if (!enif_get_tuple(env, term, &size, &terms)) {
      throw std::invalid_argument(
          "decode failed, expected shape to be a tuple");
    }

    auto vector = std::vector<int64_t>();
    vector.reserve(size);

    for (auto i = 0; i < size; i++) {
      auto elem = fine::decode<int64_t>(env, terms[i]);
      vector.push_back(elem);
    }

    return vector;
  }

  static xla::PrimitiveType decode_type(ErlNifEnv *env,
                                        const ERL_NIF_TERM &term) {
    auto [element, size] =
        fine::decode<std::tuple<fine::Atom, uint64_t>>(env, term);

    if (element == "u") {
      switch (size) {
      case 2:
        return xla::U2;
      case 4:
        return xla::U4;
      case 8:
        return xla::U8;
      case 16:
        return xla::U16;
      case 32:
        return xla::U32;
      case 64:
        return xla::U64;
      }
    }
    if (element == "s") {
      switch (size) {
      case 2:
        return xla::S2;
      case 4:
        return xla::S4;
      case 8:
        return xla::S8;
      case 16:
        return xla::S16;
      case 32:
        return xla::S32;
      case 64:
        return xla::S64;
      }
    }
    if (element == "f") {
      switch (size) {
      case 8:
        return xla::F8E5M2;
      case 16:
        return xla::F16;
      case 32:
        return xla::F32;
      case 64:
   

  static ERL_NIF_TERM encode(ErlNifEnv *env, const mlir::Type &type) {
    ERL_NIF_TERM keys[] = {
        fine::encode(env, exla::atoms::__struct__),
        fine::encode(env, exla::atoms::type),
        fine::encode(env, exla::atoms::shape),
    };

    ERL_NIF_TERM values[] = {
        fine::encode(env, exla::atoms::ElixirEXLATypespec),
        encode_type(env, type),
        encode_shape(env, type),
    };

    ERL_NIF_TERM map;
    if (!enif_make_map_from_arrays(env, keys, values, 3, &map)) {
      throw std::runtime_error("encode: failed to make a map");
    }

    return map;
  }

private:
  static ERL_NIF_TERM encode_type(ErlNifEnv *env, const mlir::Type &type) {
    if (mlir::isa<mlir::stablehlo::TokenType>(type)) {
      return fine::encode(env, exla::atoms::token);
    }

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

    if (mlir::isa<mlir::RankedTensorType>(type)) {
      auto tensor_type = mlir::cast<mlir::RankedTensorType>(type);
      auto element_type = tensor_type.getElementType();

      if (element_type.isSignlessInteger(1)) {
        type_name = exla::atoms::pred;
        type_size = 8;
      } else if (auto integer_type =
                     mlir::dyn_cast<mlir::IntegerType>(element_type)) {
        if (integer_type.isUnsigned()) {
          type_name = exla::atoms::u;
        } else {
          type_name = exla::atoms::s;
        }

        type_size = integer_type.getWidth();
      } else if (element_type.isBF16()) {
        type_name = exla::atoms::bf;
        type_size = 16;
      // Float8E4M3FNType doesn't

  static ERL_NIF_TERM encode_shape(ErlNifEnv *env, const mlir::Type &type) {
    if (mlir::isa<mlir::stablehlo::TokenType>(type)) {
      return enif_make_tuple(env, 0);
    }

    if (mlir::isa<mlir::RankedTensorType>(type)) {
      auto tensor_type = mlir::cast<mlir::RankedTensorType>(type);
      auto dims_array = tensor_type.getShape();

      auto dims = std::vector<ERL_NIF_TERM>{};
      dims.reserve(dims_array.size());

      for (auto dim : dims_array) {
        dims.push_back(fine::encode<int64_t>(env, dim));
      }

      return enif_make_tuple_from_array(env, dims.data(), dims.size());
    }

    throw std::invalid_argument("encode failed, unexpected mlir type");
  }

  static ERL_NIF_TERM encode(ErlNifEnv *env, const xla::Shape &shape) {
    ERL_NIF_TERM keys[] = {
        fine::encode(env, exla::atoms::__struct__),
        fine::encode(env, exla::atoms::type),
        fine::encode(env, exla::atoms::shape),
    };

    ERL_NIF_TERM values[] = {
        fine::encode(env, exla::atoms::ElixirEXLATypespec),
        encode_type(env, shape),
        encode_shape(env, shape),
    };

    ERL_NIF_TERM map;
    if (!enif_make_map_from_arrays(env, keys, values, 3, &map)) {
      throw std::runtime_error("encode: failed to make a map");
    }

    return map;
  }

private:
  static ERL_NIF_TERM encode_type(ErlNifEnv *env, const xla::Shape &shape) {
    auto ptype = shape.element_type();

    // Tokens are encoded as the atom :token
    if (ptype == xla::TOKEN) {
      return fine::encode(env, exla::atoms::token);
    }

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

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

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

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

  static ERL_NIF_TERM encode_shape(ErlNifEnv *env, const xla::Shape &shape) {
    auto dims = shape.dimensions();
    std::vector<ERL_NIF_TERM> dim_terms;
    dim_terms.reserve(dims.size());

    for (int64_t dim : dims) {
      dim_terms.push_back(fine::encode(env, dim));
    }

    return enif_make_tuple_from_array(env, dim_terms.data(), dim_terms.size());
  }

int get_ipc_handle(const char* memname, size_t memsize, mode_t mode) {
  int fd = shm_open(memname, O_CREAT | O_RDWR, mode);
  if (fd == -1) {
    return -1;
  }

  if (ftruncate(fd, memsize) == -1) {
    close(fd);
    return -1;
  }

  return fd;
}

int open_existing_ipc_handle(const char* memname, int* out_writable) {
  int fd = shm_open(memname, O_RDWR, 0);
  if (fd != -1) {
    *out_writable = 1;
    return fd;
  }
  if (errno == EACCES) {
    fd = shm_open(memname, O_RDONLY, 0);
    if (fd != -1) {
      *out_writable = 0;
      return fd;
    }
  }
  return -1;
}

void* open_ipc_handle(int fd, size_t memsize, int writable) {
  int prot = writable ? (PROT_READ | PROT_WRITE) : PROT_READ;
  void* ptr = mmap(NULL, memsize, prot, MAP_SHARED, fd, 0);
  if (ptr == MAP_FAILED) {
    perror("mmap");
    return nullptr;
  }
  return ptr;
}

int close_ipc_handle(int fd, void* ptr, const char* memname, size_t memsize) {
  if (munmap(ptr, memsize) == -1) {
    return -1;
  }

  if (close(fd) == -1) {
    return -1;
  }

  shm_unlink(memname);

  return 0;
}

int close_imported_ipc_handle(int fd, void* ptr, size_t memsize) {
  munmap(ptr, memsize);
  close(fd);
  return 0;
}

ERL_NIF_TERM error(ErlNifEnv *env, const char *msg)
{
  ERL_NIF_TERM atom = enif_make_atom(env, "error");
  ERL_NIF_TERM msg_term = enif_make_string(env, msg, ERL_NIF_LATIN1);
  return enif_make_tuple2(env, atom, msg_term);
}

ERL_NIF_TERM ok(ErlNifEnv *env)
{
  return enif_make_atom(env, "ok");
}

int upgrade(ErlNifEnv *env, void **priv_data, void **old_priv_data, ERL_NIF_TERM load_info) {
  // Silence "unused var" warnings.
  (void)(env);
  (void)(priv_data);
  (void)(old_priv_data);
  (void)(load_info);

  return 0;
}

int load(ErlNifEnv *,void **,ERL_NIF_TERM) {
  return 0;
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
tensor_ok(const torch::Tensor &tensor) {
  return fine::Ok(fine::make_resource<TorchTensor>(tensor));
}

c10::IntArrayRef vec_to_array_ref(const std::vector<int64_t> &vec) {
  return c10::IntArrayRef(vec);
}

torch::Device
tuple_to_device(const std::tuple<int64_t, int64_t> &device_tuple) {
  return torch::Device(
      static_cast<torch::DeviceType>(std::get<0>(device_tuple)),
      static_cast<torch::DeviceIndex>(std::get<1>(device_tuple)));
}

uint64_t elem_count(const std::vector<int64_t> &shape) {
  return std::accumulate(shape.begin(), shape.end(), 1ULL, std::multiplies<>{});
}

fine::Atom delete_tensor(ErlNifEnv *env,
                         fine::ResourcePtr<TorchTensor> tensor) {
  if (tensor->deallocate()) {
    return fine::Atom("ok");
  } else {
    // Throw exception so backend can catch and return :already_deallocated
    throw std::invalid_argument("Tensor has been deallocated");
  }
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
from_blob(ErlNifEnv *env, ErlNifBinary blob, std::vector<int64_t> shape,
          fine::Atom type_atom, std::tuple<int64_t, int64_t> device_tuple) {

  auto type = string2type(type_atom.to_string());
  auto device = tuple_to_device(device_tuple);

  // Check if binary is large enough
  if (blob.size / dtype_sizes[type_atom.to_string()] < elem_count(shape)) {
    throw std::invalid_argument(
        "Binary size is too small for the requested shape");
  }

  auto tensor = torch::from_blob(blob.data, vec_to_array_ref(shape),
                                 torch::device(torch::kCPU).dtype(type));

  if (device.type() == torch::kCPU) {
    return tensor_ok(tensor.clone());
  } else {
    return tensor_ok(tensor.to(device));
  }
}

fine::Ok<ErlNifBinary> to_blob_1(ErlNifEnv *env,
                                 fine::ResourcePtr<TorchTensor> tensor_res) {
  auto &t = get_tensor(tensor_res);
  size_t byte_size = t.nbytes();

  torch::optional<torch::Device> device = torch::device_of(t);
  torch::Tensor reshaped = t.flatten();
  void *data_ptr = reshaped.data_ptr();

  ErlNifBinary result;
  enif_alloc_binary(byte_size, &result);

  // Always copy data to avoid use-after-free when tensor is deallocated
  if (device.has_value() && device.value().type() == torch::kCPU) {
    memcpy(result.data, data_ptr, byte_size);
  } else {
    memcpy(result.data, reshaped.to(torch::kCPU).data_ptr(), byte_size);
  }

  return fine::Ok(result);
}

fine::Ok<ErlNifBinary> to_blob_2(ErlNifEnv *env,
                                 fine::ResourcePtr<TorchTensor> tensor_res,
                                 int64_t limit) {
  auto &t = get_tensor(tensor_res);
  size_t byte_size = limit * t.itemsize();

  torch::optional<torch::Device> device = torch::device_of(t);
  torch::Tensor reshaped =
      (byte_size < t.nbytes()) ? t.flatten().slice(0, 0, limit) : t.flatten();
  void *data_ptr = reshaped.data_ptr();

  ErlNifBinary result;
  enif_alloc_binary(byte_size, &result);

  // Always copy data to avoid use-after-free when tensor is deallocated
  if (device.has_value() && device.value().type() == torch::kCPU) {
    memcpy(result.data, data_ptr, byte_size);
  } else {
    memcpy(result.data, reshaped.to(torch::kCPU).data_ptr(), byte_size);
  }

  return fine::Ok(result);
}

fine::Ok<torch::Scalar> item(ErlNifEnv *env,
                             fine::ResourcePtr<TorchTensor> tensor) {
  return fine::Ok(get_tensor(tensor).item());
}

fine::Ok<fine::Atom> scalar_type(ErlNifEnv *env,
                                 fine::ResourcePtr<TorchTensor> tensor) {
  const std::string *type_name = type2string(get_tensor(tensor).scalar_type());
  if (type_name != nullptr) {
    return fine::Ok(fine::Atom(*type_name));
  } else {
    throw std::runtime_error("Could not determine tensor type.");
  }
}

fine::Ok<fine::Term> shape(ErlNifEnv *env,
                           fine::ResourcePtr<TorchTensor> tensor) {
  auto &t = get_tensor(tensor);
  std::vector<ERL_NIF_TERM> sizes;
  for (int64_t dim = 0; dim < t.dim(); dim++) {
    sizes.push_back(fine::encode(env, t.size(dim)));
  }
  // Return as tuple (not list) since Elixir expects {} not []
  return fine::Ok(
      fine::Term(enif_make_tuple_from_array(env, sizes.data(), sizes.size())));
}

bool mps_is_available(ErlNifEnv *env) {
#ifdef MAC_ARM64
  return at::hasMPS();
#else
  return false;
#endif
}

int64_t cuda_device_count(ErlNifEnv *env) {
  return static_cast<int64_t>(torch::cuda::device_count());
}

fine::Ok<int64_t> nbytes(ErlNifEnv *env,
                         fine::ResourcePtr<TorchTensor> tensor) {
  return fine::Ok(static_cast<int64_t>(get_tensor(tensor).nbytes()));
}

fine::Ok<std::vector<fine::ResourcePtr<TorchTensor>>>
split(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
      int64_t batch_size) {
  auto tensors = torch::split(get_tensor(tensor), batch_size);
  std::vector<fine::ResourcePtr<TorchTensor>> results;
  for (const auto &t : tensors) {
    results.push_back(fine::make_resource<TorchTensor>(t));
  }
  return fine::Ok(results);
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
reshape(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
        std::vector<int64_t> shape) {
  return tensor_ok(torch::reshape(get_tensor(tensor), vec_to_array_ref(shape)));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
to_type(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
        fine::Atom type_atom) {
  auto type = string2type(type_atom.to_string());
  return tensor_ok(get_tensor(tensor).toType(type));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
to_device(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
          std::tuple<int64_t, int64_t> device_tuple) {
  TORCH_CATCH_ERROR(
      [&]() {
        auto device = tuple_to_device(device_tuple);
        return tensor_ok(get_tensor(tensor).to(device));
      }(),
      "Device transfer");
}

  TORCH_CATCH_ERROR(
      [&]() {
        auto device = tuple_to_device(device_tuple);
        return tensor_ok(get_tensor(tensor).to(device));
      }(),

std::variant<fine::Ok<fine::ResourcePtr<TorchTensor>>>
squeeze(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
        std::optional<int64_t> dim) {
  if (dim.has_value()) {
    return tensor_ok(torch::squeeze(get_tensor(tensor), dim.value()));
  } else {
    return tensor_ok(torch::squeeze(get_tensor(tensor)));
  }
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
broadcast_to(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor,
             std::vector<int64_t> shape) {
  return tensor_ok(
      torch::broadcast_to(get_tensor(tensor), vec_to_array_ref(shape)).clone());
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
transpose(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> tensor, int64_t dim0,
          int64_t dim1) {
  return tensor_ok(torch::transpose(get_tensor(tensor), dim0, dim1));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
slice(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> input,
      std::vector<int64_t> starts, std::vector<int64_t> lengths,
      std::vector<int64_t> strides) {

  torch::Tensor result = get_tensor(input);
  auto shape = result.sizes();

  for (size_t dim = 0; dim < starts.size(); dim++) {
    int64_t start = starts[dim];
    int64_t stride = strides[dim];
    int64_t length = lengths[dim];
    int64_t end = std::min(start + length, shape[dim]);

    result = result.slice(dim, start, end, stride);
  }

  // Clone the result to ensure memory ownership
  return tensor_ok(result.clone());
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
concatenate(ErlNifEnv *env,
            std::vector<fine::ResourcePtr<TorchTensor>> tensor_list,
            int64_t dim) {
  std::vector<torch::Tensor> tensors;
  for (const auto &t : tensor_list) {
    tensors.push_back(get_tensor(t));
  }
  return tensor_ok(torch::cat(tensors, dim));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
gather(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> input,
       fine::ResourcePtr<TorchTensor> index, int64_t dim) {
  return tensor_ok(torch::gather(get_tensor(input), dim, get_tensor(index)));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
index_put(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> input,
          std::vector<fine::ResourcePtr<TorchTensor>> indices,
          fine::ResourcePtr<TorchTensor> values, bool accumulate) {
  TORCH_CATCH_ERROR(
      [&]() {
        c10::List<std::optional<at::Tensor>> torch_indices;
        for (const auto &idx : indices) {
          torch_indices.push_back(get_tensor(idx));
        }

        torch::Tensor result = get_tensor(input).clone();
        result.index_put_(torch_indices, get_tensor(values), accumulate);
        return tensor_ok(result);
      }(),
      "index_put");
}

  TORCH_CATCH_ERROR(
      [&]() {
        c10::List<std::optional<at::Tensor>> torch_indices;
        for (const auto &idx : indices) {
          torch_indices.push_back(get_tensor(idx));
        }

        torch::Tensor result = get_tensor(input).clone();
        result.index_put_(torch_indices, get_tensor(values), accumulate);
        return tensor_ok(result);
      }(),

fine::Ok<fine::ResourcePtr<TorchTensor>>
index(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> input,
      std::vector<fine::ResourcePtr<TorchTensor>> indices) {

  c10::List<std::optional<at::Tensor>> torch_indices;
  for (const auto &idx : indices) {
    torch_indices.push_back(get_tensor(idx));
  }

  return tensor_ok(get_tensor(input).index(torch_indices));
}

fine::Ok<fine::ResourcePtr<TorchTensor>>
argsort(ErlNifEnv *env, fine::ResourcePtr<TorchTensor> input, bool stable,
        int64_t dim, bool descending) {
  return tensor_ok(torch::argsort(get_tensor(input), stable, dim, descending));
}