Function eval_impl

mlx/transforms.cpp:74–314

int detail::RetainGraph::tracing_counter{0};

array eval_impl(std::vector<array> outputs, bool async) {
  std::deque<array> tape;

  // Make an effort to choose a good output stream
  Stream stream = default_stream(default_device());
  for (auto& o : outputs) {
    if (o.status() == array::Status::unscheduled && o.has_primitive()) {
      stream = o.primitive().stream();
      break;
    }
  }

  // Map of array id that needs fence and stream it's computed on
  std::unordered_map<uintptr_t, std::pair<uint32_t, bool>> needs_fence;

  auto synchronizer = array(
      {}, bool_, std::make_shared<Synchronizer>(stream), std::move(outputs));

  // Stream fences for inter-stream synchronization
  std::unordered_map<uint32_t, Fence> fences;

  // Stream events for synchronization after eval
  std::unordered_map<uint32_t, Event> events;
  {
    auto e = Event{stream};
    e.set_value(1);
    synchronizer.attach_event(e);
    events.emplace(stream.index, std::move(e));
  }

  {
    // Record the degree of each input
    std::unordered_map<std::uintptr_t, int> cache;

    std::stack<std::pair<std::reference_wrapper<array>, int>> dfs;
    dfs.emplace(synchronizer, 0);
    while (!dfs.empty()) {
      auto& [a_ref, idx] = dfs.top();
      auto& a = a_ref.get();

      if (idx < a.inputs().size()) {
        // Add an input, and continue
        auto& in = a.inputs()[idx++];

        if (in.status() == array::Status::unscheduled) {
          if (async && in.is_tracer()) {
            throw std::invalid_argument(
                "[async_eval] Not allowed inside a graph transformation.");
          }
          if (!in.has_primitive()) {
            if (in.is_tracer()) {
              throw std::invalid_argument(
                  "[eval] Attempting to eval an array during function"
                  " transformations like compile or vmap is not allowed.");
            }
            throw std::runtime_error(
                "[eval] Attempting to eval an array without a primitive.\n"
                "If you are compiling a function, make sure all the inputs "