/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file cuda_module.cc
 * \brief CUDAModuleNode — runtime-side, plugin-only.  Reachable from C++ only
 *        through the FFI registry keys "ffi.Module.create.cuda" and
 *        "ffi.Module.load_from_bytes.cuda".  No exported header — codegen-side
 *        construction goes through src/target/cuda/cuda_fallback_module.h.
 */
#include <cuda.h>
#include <cuda_runtime.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/cuda/base.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>

#include <array>
#include <mutex>
#include <string>
#include <vector>

#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
#include "../../../runtime/thread_storage_scope.h"
#include "../../../support/bytes_io.h"

namespace tvm {
namespace runtime {

#ifndef CUDA_DRIVER_CALL
#define CUDA_DRIVER_CALL(x)                                             \
  {                                                                     \
    CUresult result = x;                                                \
    if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) { \
      const char* msg;                                                  \
      cuGetErrorName(result, &msg);                                     \
      TVM_FFI_THROW(CUDAError) << "" #x " failed with error: " << msg;  \
    }                                                                   \
  }
#endif

// Maximum number of GPU supported in CUDAModule (file-local).
static constexpr const int kMaxNumGPUs = 32;

// Module to support thread-safe multi-GPU execution.
// cuModule is a per-GPU module
// The runtime will contain a per-device module table
// The modules will be lazily loaded
class CUDAModuleNode : public ffi::ModuleObj {
 public:
  CUDAModuleNode(ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, FunctionInfo> fmap,
                 ffi::Map<ffi::String, ffi::String> source)
      : code_(code), fmt_(fmt), fmap_(fmap), source_(source) {
    std::fill(module_.begin(), module_.end(), nullptr);
  }
  // destructor
  ~CUDAModuleNode() {
    for (size_t i = 0; i < module_.size(); ++i) {
      if (module_[i] != nullptr) {
        cudaError_t set_err = cudaSetDevice(static_cast<int>(i));
        if (set_err != cudaSuccess && set_err != cudaErrorCudartUnloading) {
          continue;
        }
        CUresult result = cuModuleUnload(module_[i]);
        // Ignore errors during cleanup - context may be shutting down
        (void)result;
      }
    }
  }

  const char* kind() const final { return "cuda"; }

  /*! \brief Get the property of the runtime module .*/
  int GetPropertyMask() const final {
    return ffi::Module::kBinarySerializable | ffi::Module::kRunnable;
  }

  ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) final;

  ffi::Bytes SaveToBytes() const final {
    // Format: [fmt][fmap][code].  Source map is in-memory inspection only and
    // is NEVER serialized — it is lost on save/load round-trip (matches
    // upstream behavior; the receiver rebuilds source from code bytes if
    // possible).  CUDAFallbackModuleNode::SaveToBytes (in
    // src/target/cuda/cuda_fallback_module.cc) MUST mirror this format
    // byte-for-byte; see one-way comment there.
    std::string buffer;
    support::BytesOutStream stream(&buffer);
    stream.Write(fmt_);
    stream.Write(fmap_);
    stream.Write(code_);
    return ffi::Bytes(std::move(buffer));
  }

  ffi::String InspectSource(const ffi::String& format) const final {
    // For known compiled formats, return code as string when format matches.
    if (format == fmt_) {
      return ffi::String(code_.data(), code_.size());
    }
    // Look up the source map for an exact match (e.g. "cuda" returns the
    // original C++ source, populated by codegen at construction time).
    if (auto it = source_.find(format); it != source_.end()) {
      return (*it).second;
    }
    // Empty-format (`mod.get_source()`) — prefer the `cuda` source if present,
    // else fall back to code-as-string when fmt_ is textual.
    if (format.empty()) {
      if (auto it = source_.find("cuda"); it != source_.end()) {
        return (*it).second;
      }
      if (fmt_ == "ptx" || fmt_ == "cuda") {
        return ffi::String(code_.data(), code_.size());
      }
    }
    return ffi::String();
  }

  // get a CUfunction from primary context in device_id
  CUfunction GetFunc(int device_id, const std::string& func_name) {
    std::lock_guard<std::mutex> lock(mutex_);
    // must recheck under the lock scope
    if (module_[device_id] == nullptr) {
      CUDA_DRIVER_CALL(cuModuleLoadData(&(module_[device_id]), code_.data()));
      static auto nvshmem_init_hook = ffi::Function::GetGlobal("runtime.nvshmem.cumodule_init");
      if (nvshmem_init_hook.has_value()) {
        (*nvshmem_init_hook)(static_cast<void*>(module_[device_id]));
      }
    }
    CUfunction func;
    CUresult result = cuModuleGetFunction(&func, module_[device_id], func_name.c_str());
    if (result != CUDA_SUCCESS) {
      const char* msg;
      cuGetErrorName(result, &msg);
      TVM_FFI_THROW(CUDAError) << "cuModuleGetFunction " << func_name
                               << " failed with error: " << msg;
    }
    return func;
  }

  /*!
   * \brief JIT-compile raw CUDA C++ source to PTX/cubin/fatbin via the Python
   *        compile callback.  Called from BOTH the
   *        "ffi.Module.create.cuda" lambda (when the codegen hands us
   *        fmt=="cuda") AND from LoadFromBytes (when the saved-on-disk fmt is
   *        "cuda" — the cross-compile receiver path).
   *
   * \param source Raw CUDA C++ source (text).
   * \return Compiled binary bytes.  Determination of the compiled format
   *         (ptx vs cubin vs fatbin) is left to the caller (heuristic on
   *         first byte: '/' → ptx-text, otherwise binary).
   */
  static ffi::Bytes JitCompileFromSource(const ffi::String& source) {
    // Registry: "tvm_callback_cuda_compile" — Python-side nvcc/nvrtc wrapper.
    // Grep hint: grep -rn 'tvm_callback_cuda_compile' src/ python/
    auto fcompile = ffi::Function::GetGlobal("tvm_callback_cuda_compile");
    TVM_FFI_CHECK(fcompile.has_value(), RuntimeError)
        << "fmt=='cuda' requires tvm_callback_cuda_compile to be registered. "
        << "Import tvm.support.nvcc.";
    return (*fcompile)(source).cast<ffi::Bytes>();
  }

  /*! \brief Pick the compiled format from the JIT output's first byte. */
  static ffi::String DetermineCompiledFormat(const ffi::Bytes& compiled) {
    if (compiled.size() > 0 && compiled.data()[0] == '/') {
      return ffi::String("ptx");
    }
    return ffi::String("cubin");
  }

 private:
  // The binary data (compiled PTX/cubin/fatbin, or raw CUDA source if fmt == "cuda").
  ffi::Bytes code_;
  // The format of code_.
  ffi::String fmt_;
  // function information table.
  ffi::Map<ffi::String, FunctionInfo> fmap_;
  // In-memory source map for InspectSource — never serialized.
  ffi::Map<ffi::String, ffi::String> source_;
  // the internal modules per GPU, to be lazily initialized.
  std::array<CUmodule, kMaxNumGPUs> module_;
  // internal mutex when updating the module
  std::mutex mutex_;
};

// a wrapped function class to get packed func.
class CUDAWrappedFunc {
 public:
  // initialize the CUDA function.
  void Init(CUDAModuleNode* m, ffi::ObjectPtr<ffi::Object> sptr, const std::string& func_name,
            size_t num_void_args, const ffi::Array<ffi::String>& launch_param_tags) {
    m_ = m;
    sptr_ = sptr;
    func_name_ = func_name;
    std::fill(fcache_.begin(), fcache_.end(), nullptr);
    launch_param_config_.Init(num_void_args, launch_param_tags);
  }
  // invoke the function with void arguments
  void operator()(ffi::PackedArgs args, ffi::Any* rv, void** void_args) const {
    int device_id;
    TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&device_id));
    ThreadWorkLoad wl = launch_param_config_.Extract(args);

    if (fcache_[device_id] == nullptr) {
      fcache_[device_id] = m_->GetFunc(device_id, func_name_);
      if (wl.dyn_shmem_size >= (48 << 10)) {
        // Assumption: dyn_shmem_size doesn't change across different invocations of
        // fcache_[device_id]
        CUresult result = cuFuncSetAttribute(
            fcache_[device_id], CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, wl.dyn_shmem_size);
        if (result != CUDA_SUCCESS) {
          TVM_FFI_THROW(InternalError)
              << "Failed to set the allowed dynamic shared memory size to " << wl.dyn_shmem_size;
        }
      }
    }
    CUstream strm = static_cast<CUstream>(TVMFFIEnvGetStream(kDLCUDA, device_id));
    std::vector<CUlaunchAttribute> attrs;

    // 1) Cluster
    if (wl.cluster_dim(0) != 1 || wl.cluster_dim(1) != 1 || wl.cluster_dim(2) != 1) {
      CUlaunchAttribute attr{};
      attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
      attr.value.clusterDim.x = wl.cluster_dim(0);
      attr.value.clusterDim.y = wl.cluster_dim(1);
      attr.value.clusterDim.z = wl.cluster_dim(2);
      attrs.push_back(attr);
    }

    // 1b) Preferred cluster (CUDA 12.8+, cudaLaunchAttributePreferredClusterDimension)
    if (wl.preferred_cluster_dim(0) != 1 || wl.preferred_cluster_dim(1) != 1 ||
        wl.preferred_cluster_dim(2) != 1) {
      CUlaunchAttribute attr{};
      attr.id = CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION;
      attr.value.clusterDim.x = wl.preferred_cluster_dim(0);
      attr.value.clusterDim.y = wl.preferred_cluster_dim(1);
      attr.value.clusterDim.z = wl.preferred_cluster_dim(2);
      attrs.push_back(attr);
    }

    // 2) Programmatic stream serialization
    if (launch_param_config_.use_programtic_dependent_launch()) {
      CUlaunchAttribute attr{};
      attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION;
      attr.value.programmaticStreamSerializationAllowed = 1;
      attrs.push_back(attr);
    }

    // 3) Cooperative
    if (launch_param_config_.use_cooperative_launch()) {
      CUlaunchAttribute attr{};
      attr.id = CU_LAUNCH_ATTRIBUTE_COOPERATIVE;
      attr.value.cooperative = 1;
      attrs.push_back(attr);
    }

    // 4) Launch
    CUlaunchConfig config{};
    config.gridDimX = wl.grid_dim(0);
    config.gridDimY = wl.grid_dim(1);
    config.gridDimZ = wl.grid_dim(2);
    config.blockDimX = wl.block_dim(0);
    config.blockDimY = wl.block_dim(1);
    config.blockDimZ = wl.block_dim(2);
    config.sharedMemBytes = wl.dyn_shmem_size;
    config.hStream = strm;
    config.attrs = attrs.empty() ? nullptr : attrs.data();
    config.numAttrs = static_cast<unsigned int>(attrs.size());

    CUresult result = cuLaunchKernelEx(&config, fcache_[device_id], void_args, nullptr);

    if (result != CUDA_SUCCESS && result != CUDA_ERROR_DEINITIALIZED) {
      const char* msg;
      cuGetErrorName(result, &msg);
      std::ostringstream os;
      os << "CUDALaunch Error: " << msg << "\n"
         << " grid=(" << wl.grid_dim(0) << "," << wl.grid_dim(1) << "," << wl.grid_dim(2) << "), "
         << " block=(" << wl.block_dim(0) << "," << wl.block_dim(1) << "," << wl.block_dim(2)
         << ")\n";
      ffi::String cuda = m_->InspectSource("");
      if (cuda.length() != 0) {
        os << "// func_name=" << func_name_ << "\n"
           << "// CUDA Source\n"
           << "// -----------\n"
           << cuda;
      }
      TVM_FFI_THROW(InternalError) << os.str();
    }
  }

 private:
  // internal module
  CUDAModuleNode* m_;
  // the resource holder
  ffi::ObjectPtr<ffi::Object> sptr_;
  // The name of the function.
  std::string func_name_;
  // Device function cache per device.
  // mark as mutable, to enable lazy initialization
  mutable std::array<CUfunction, kMaxNumGPUs> fcache_;
  // launch parameters configuration
  LaunchParamConfig launch_param_config_;
};

ffi::Optional<ffi::Function> CUDAModuleNode::GetFunction(const ffi::String& name) {
  ffi::ObjectPtr<ffi::Object> sptr_to_self = ffi::GetObjectPtr<ffi::Object>(this);
  TVM_FFI_ICHECK_EQ(sptr_to_self.get(), this);
  auto opt_info = fmap_.Get(name);
  if (!opt_info.has_value()) return ffi::Function();
  FunctionInfo info = opt_info.value();
  CUDAWrappedFunc f;
  f.Init(this, sptr_to_self, name, info->arg_types.size(), info->launch_param_tags);
  return PackFuncVoidAddr(f, info->arg_types, info->arg_extra_tags);
}

// Construct a CUDAModuleNode from in-memory payload.  When fmt == "cuda" the
// code is raw CUDA C++ source — JIT-compile via the Python callback, then
// re-tag with the resulting compiled format ("ptx" / "cubin").
static ffi::Module CUDAModuleCreateImpl(ffi::Bytes code, ffi::String fmt,
                                        ffi::Map<ffi::String, FunctionInfo> fmap,
                                        ffi::Map<ffi::String, ffi::String> source) {
  if (fmt == "cuda") {
    // Stash the CUDA source for InspectSource before we replace `code` with
    // the JIT output.
    if (source.find("cuda") == source.end()) {
      source.Set("cuda", ffi::String(code.data(), code.size()));
    }
    ffi::Bytes compiled =
        CUDAModuleNode::JitCompileFromSource(ffi::String(code.data(), code.size()));
    fmt = CUDAModuleNode::DetermineCompiledFormat(compiled);
    code = std::move(compiled);
  }
  auto n = ffi::make_object<CUDAModuleNode>(code, fmt, fmap, source);
  return ffi::Module(n);
}

static ffi::Module CUDAModuleLoadFromBytes(const ffi::Bytes& bytes) {
  support::BytesInStream stream(bytes);
  ffi::String fmt;
  ffi::Map<ffi::String, FunctionInfo> fmap;
  ffi::Bytes code;
  stream.Read(&fmt);
  TVM_FFI_ICHECK(stream.Read(&fmap));
  stream.Read(&code);
  // Source map is not serialized — it is lost on save/load round-trip.
  // If the receiver wants InspectSource("cuda") to work, the saved bytes must
  // have been written with fmt=="cuda" so the JIT path below re-stuffs the
  // source map with the original C++ source.
  return CUDAModuleCreateImpl(std::move(code), std::move(fmt), std::move(fmap),
                              ffi::Map<ffi::String, ffi::String>());
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  // Registry: "ffi.Module.create.cuda" — codegen-time CUDA module factory.
  // Used by src/target/cuda/cuda_fallback_module.h:CUDAModuleCreateWithFallback.
  // Registry: "ffi.Module.load_from_bytes.cuda" — disk loader.  Only this
  // (real) module registers a loader; the fallback module is codegen-time only.
  refl::GlobalDef()
      .def("ffi.Module.load_from_bytes.cuda", CUDAModuleLoadFromBytes)
      .def("ffi.Module.create.cuda",
           [](ffi::Bytes code, ffi::String fmt, ffi::Map<ffi::String, FunctionInfo> fmap,
              ffi::Map<ffi::String, ffi::String> source) {
             return CUDAModuleCreateImpl(std::move(code), std::move(fmt), std::move(fmap),
                                         std::move(source));
           });
}
}  // namespace runtime
}  // namespace tvm
