/*!
 * \file builtin.cc
 * \brief Multi-GPU builtin functions in MLC LLM.
 */
#ifndef MLC_SINGLE_GPU_ONLY

#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/vm/vm.h>

namespace mlc {
namespace llm {
namespace multi_gpu {

using namespace tvm::runtime;
using tvm::ffi::Array;
using tvm::ffi::ObjectRef;
using tvm::ffi::Optional;
using tvm::ffi::Shape;

ObjectRef DispatchFunctionByGroup(tvm::ffi::AnyView vm_arg,
                                  Array<Array<ObjectRef>> funcs_and_args) {
  using namespace vm;
  VirtualMachine* vm = VirtualMachine::GetContextPtr(vm_arg);
  DiscoWorker* worker = DiscoWorker::ThreadLocal();
  int world_size = worker->num_workers;
  int group_size = worker->num_workers / worker->num_groups;
  int num_group = world_size / group_size;
  TVM_FFI_ICHECK_EQ(funcs_and_args.size(), num_group)
      << "Number of groups mismatches. There are " << num_group
      << " groups while the function/arg array has " << funcs_and_args.size() << " elements.";

  int group_id = worker->worker_id / group_size;
  TVM_FFI_ICHECK(!funcs_and_args[group_id].empty())
      << "No function is provided for group " << group_id;
  VMClosure func = funcs_and_args[group_id][0].as_or_throw<VMClosure>();

  int num_args = static_cast<int>(funcs_and_args[group_id].size()) - 1;
  std::vector<tvm::ffi::AnyView> packed_args(num_args);
  for (int i = 0; i < num_args; ++i) {
    // NOTE: Need explicily define `arg` so that the argument does not
    // have type code kTVMObjectRValueRefArg.
    packed_args[i] = funcs_and_args[group_id][1 + i];
  }

  tvm::ffi::Any rv;
  vm->InvokeClosurePacked(funcs_and_args[group_id][0].as_or_throw<VMClosure>(),
                          tvm::ffi::PackedArgs(packed_args.data(), packed_args.size()), &rv);
  return rv.cast<ObjectRef>();
}

ObjectRef SendFromLastGroupToWorker0(Tensor send, Optional<Tensor> recv, Shape shape,
                                     DLDataType dtype) {
  DiscoWorker* worker = DiscoWorker::ThreadLocal();
  int worker_id = worker->worker_id;
  int world_size = worker->num_workers;
  int group_size = worker->num_workers / worker->num_groups;
  TVM_FFI_ICHECK_NE(world_size, group_size) << "Cannot perform when there is only one group.";
  int sender_id = world_size - group_size;
  if (worker_id == 0) {
    TVM_FFI_ICHECK(recv.has_value()) << "The receive Tensor is undefined for worker 0.";
    Tensor recv_arr = recv.value().CreateView(shape, dtype);
    RecvFromWorker(recv_arr, sender_id);
    return recv_arr;
  } else if (worker_id == sender_id) {
    TVM_FFI_ICHECK_EQ(send->dtype, dtype)
        << "The src Tensor has mismatched dtype than the expected dtype.";
    TVM_FFI_ICHECK_EQ(send->ndim, shape.size())
        << "The src Tensor has mismatched shape than the expected shape.";
    for (int i = 0; i < send->ndim; ++i) {
      TVM_FFI_ICHECK_EQ(send->shape[i], shape[i])
          << "The src Tensor has mismatched shape than the expected shape.";
    }
    SendToWorker(send, /*receiver_id=*/0);
    return recv.value_or(Tensor(nullptr));
  }

  // We only process for worker 0 and the first worker of the last group.
  // For other workers, we return the input object.
  return recv.value_or(Tensor(nullptr));
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef()
      .def("mlc.multi_gpu.DispatchFunctionByGroup", DispatchFunctionByGroup)
      .def("mlc.multi_gpu.SendFromLastGroupToWorker0", SendFromLastGroupToWorker0);
}

}  // namespace multi_gpu
}  // namespace llm
}  // namespace mlc

#endif  // MLC_SINGLE_GPU_ONLY