/*************************************************************************
 * Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * See LICENSE.txt for license information
 ************************************************************************/

#include "nccl.h"
#include "debug.h"
#include "rocmwrap.h"
#include "hsa/hsa.h"
#include "param.h"
#include "bootstrap.h"

#include <unistd.h>
#include <sys/utsname.h>
#include <fstream>

#define DECLARE_ROCM_PFN(symbol) PFN_##symbol pfn_##symbol = nullptr

// DMA-BUF feature gate: stays NULL when the platform does not support DMA-BUF.
// hsa_init/hsa_system_get_info/hsa_status_string are called directly via the
// hsa-runtime64 library librccl links against (no dlopen/dlsym).
DECLARE_ROCM_PFN(hsa_amd_portable_export_dmabuf);
NCCL_PARAM(DmaBufEnable, "DMABUF_ENABLE", 1);
RCCL_PARAM(ForceEnableDMABUF, "FORCE_ENABLE_DMABUF", 0);

static uint16_t version_major, version_minor;

int ncclCudaDriverVersionCache = -1;
bool ncclCudaLaunchBlocking = false;

static pthread_once_t initOnceControl = PTHREAD_ONCE_INIT;
static ncclResult_t initResult;

// This env var (NCCL_CUMEM_ENABLE) toggles cuMem API usage
NCCL_PARAM(CuMemEnable, "CUMEM_ENABLE", 0);
NCCL_PARAM(CuMemHostEnable, "CUMEM_HOST_ENABLE", -1);
// Handle type used for cuMemCreate()
CUmemAllocationHandleType ncclCuMemHandleType = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;

static int ncclCuMemSupported = 0;

// cuMem VMM API availability by HIP/driver version.
#define NCCL_CUMEM_NATIVE_MIN_VERSION   71260540
#define NCCL_CUMEM_BACKPORT_MIN_VERSION 70051831
#define NCCL_CUMEM_BACKPORT_MAX_VERSION 70060000

#define NCCL_CUMEM_VERSION_SUPPORTED(version)                  \
  ((version) >= NCCL_CUMEM_NATIVE_MIN_VERSION ||               \
   ((version) >= NCCL_CUMEM_BACKPORT_MIN_VERSION &&            \
    (version) < NCCL_CUMEM_BACKPORT_MAX_VERSION))

#define KERNEL_VERSION_CODE(major, minor) ((major << 16) | (minor << 8))

static int ncclGetKernelVersionCode() {
  struct utsname u;
  int major = 0, minor = 0;

  if (uname(&u) != 0) return -1;
  sscanf(u.release, "%d.%d", &major, &minor);
  INFO(NCCL_INIT, "Kernel version %d.%d", major, minor);

  return KERNEL_VERSION_CODE(major, minor);
}

// Determine whether CUMEM & VMM RDMA is supported on this platform
int ncclIsCuMemSupported() {
  CUdevice currentDev;
  int cudaDev;
  int cudaDriverVersion;
  int flag = 0;
  int supported = 1;
  ncclResult_t ret = ncclSuccess;

  if (ncclGetKernelVersionCode() < KERNEL_VERSION_CODE(6, 8)) {
    WARN("cuMem support requires Linux kernel >= 6.8");
    supported = 0;
  }
  CUDACHECKGOTO(cudaDriverGetVersion(&cudaDriverVersion), ret, error);
  {
    // Block scope prevents the goto in CUDACHECKGOTO from jumping over the bool initialization.
    bool cuMemSupported = NCCL_CUMEM_VERSION_SUPPORTED(cudaDriverVersion);
    if (!cuMemSupported) {
      WARN("cuMem support requires HIP_VERSION >= 7.12.60540 (or ROCm 7.0.2.x backport)");
      supported = 0;
    }
  }
  CUDACHECKGOTO(cudaGetDevice(&cudaDev), ret, error);
  if (CUPFN(cuMemCreate) == NULL) supported = 0;
  CUCHECKGOTO(cuDeviceGet(&currentDev, cudaDev), ret, error);
  // Query device to see if CUMEM VMM support is available
  CUCHECKGOTO(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, currentDev), ret, error);
  if (!flag) {
    WARN("cuMem support requires VMM RDMA support");
    supported = 0;
  }

  return supported;
error:
  return (ret == ncclSuccess);
}

int ncclCuMemEnable() {
#if NCCL_CUMEM_VERSION_SUPPORTED(HIP_VERSION)
  int param = ncclParamCuMemEnable();
  return param >= 0 ? param : (param == -2 && ncclCuMemSupported);
#else
  if (ncclParamCuMemEnable() > 0)
    WARN("NCCL_CUMEM_ENABLE=1 is set but cuMem VMM APIs are unavailable in this build (HIP_VERSION=%d); disabling cuMem", HIP_VERSION);
  return 0;
#endif
}

static int ncclCumemHostEnable = -1;
int ncclCuMemHostEnable() {
  if (ncclCumemHostEnable != -1)
    return ncclCumemHostEnable;
#if HIP_VERSION < 71260540
  ncclCumemHostEnable = 0;
  return ncclCumemHostEnable;
#else
  ncclResult_t ret = ncclSuccess;
  int cudaDriverVersion;
  int paramValue = -1;
  CUDACHECKGOTO(cudaDriverGetVersion(&cudaDriverVersion), ret, error);
  if (cudaDriverVersion < 71260540) {
    ncclCumemHostEnable = 0;
  }
  else {
    paramValue = ncclParamCuMemHostEnable();
    if (paramValue != -1)
      ncclCumemHostEnable = paramValue;
    else
      ncclCumemHostEnable = (cudaDriverVersion >= 71260540) ? 1 : 0;
    if (ncclCumemHostEnable) {
      // Verify that host allocations actually work.  Docker in particular is known to disable "get_mempolicy",
      // causing such allocations to fail (this can be fixed by invoking Docker with "--cap-add SYS_NICE").
      int cudaDev;
      CUdevice currentDev;
      int cpuNumaNodeId = -1;
      CUmemAllocationProp prop = {};
      size_t granularity = 0;
      size_t size;
      CUmemGenericAllocationHandle handle;
      CUDACHECK(cudaGetDevice(&cudaDev));
      CUCHECK(cuDeviceGet(&currentDev, cudaDev));
      CUCHECK(cuDeviceGetAttribute(&cpuNumaNodeId, hipDeviceAttributeHostNumaId, currentDev));
      if (cpuNumaNodeId < 0) cpuNumaNodeId = 0;
      // CLR rejects HostNuma; probe with Host to match alloc.h's ncclCuMemHostAlloc.
      prop.location.type = hipMemLocationTypeHost;
      prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
      prop.requestedHandleTypes = ncclCuMemHandleType;
      // HIP/CLR requires host id to be 0. cpuNumaNodeId can exceed GPU count and fail.
      prop.location.id = 0;  // ignored on the Host path
      CUCHECK(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
      size = 1;
      ALIGN_SIZE(size, granularity);
      if (CUPFN(cuMemCreate(&handle, size, &prop, 0)) != CUDA_SUCCESS) {
        INFO(NCCL_INIT, "cuMem host allocations do not appear to be working; falling back to a /dev/shm/ based "
             "implementation. This could be due to the container runtime disabling NUMA support. "
             "To disable this warning, set NCCL_CUMEM_HOST_ENABLE=0");
        ncclCumemHostEnable = 0;
      } else {
        CUCHECK(cuMemRelease(handle));
      }
    }
  }
  return ncclCumemHostEnable;
error:
  return (ret == ncclSuccess);
#endif
}

static void initOnceFunc() {
  do {
    char* val = getenv("CUDA_LAUNCH_BLOCKING");
    ncclCudaLaunchBlocking = val!=nullptr && val[0]!=0 && !(val[0]=='0' && val[1]==0);
  } while (0);

  bool dmaBufSupport = false;
  hsa_status_t res;

  /*
   * The HSA (ROCr) runtime is directly linked into librccl via
   * hsa-runtime64::hsa-runtime64; its entry points are resolved by the dynamic
   * loader through librccl's RPATH (the same one that resolves libamdhip64).
   * No dlopen/dlsym and no library-name string are needed here.
   */
  res = hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MAJOR, &version_major);
  if (res != 0) {
    WARN("hsa_system_get_info failed with %d", res);
    goto error;
  }
  res = hsa_system_get_info(HSA_SYSTEM_INFO_VERSION_MINOR, &version_minor);
  if (res != 0) {
    WARN("hsa_system_get_info failed with %d", res);
    goto error;
  }

  INFO(NCCL_INIT, "ROCr version %d.%d", version_major, version_minor);

  //if (hsaDriverVersion < ROCR_DRIVER_MIN_VERSION) {
    // WARN("ROCr Driver version found is %d. Minimum requirement is %d", hsaDriverVersion, ROCR_DRIVER_MIN_VERSION);
    // Silently ignore version check mismatch for backwards compatibility
    //goto error;
  //}

  // Determine whether we support the cuMem APIs or not
  ncclCuMemSupported = ncclIsCuMemSupported();

  /* DMA-BUF support */
  //ROCm support
  if(rcclParamForceEnableDMABUF())
  {
      dmaBufSupport = 1;
      WARN("DMA_BUF Support is force enabled, so explicitly setting RCCL_FORCE_ENABLE_DMABUF=1");
  }
  else if (ncclCuMemEnable() && ncclParamDmaBufEnable() == 0)
  {
    dmaBufSupport = 1;
    WARN("NCCL_CUMEM_ENABLE is set but NCCL_DMABUF_ENABLE is not. Forcefully enabling DMA-BUF for hipMem.");
  }
  else if (ncclParamDmaBufEnable() == 0)
  {
    INFO(NCCL_INIT, "Dmabuf feature disabled without NCCL_DMABUF_ENABLE=1");
    goto error;
  }

  // ROCr checks
  res = hsa_system_get_info((hsa_system_info_t) 0x204, &dmaBufSupport);
  if (res != HSA_STATUS_SUCCESS || !dmaBufSupport){
    INFO(NCCL_INIT, "Current version of ROCm does not support dmabuf feature.");
    goto error;
  }
  else if (hsa_amd_portable_export_dmabuf == nullptr) {
    // The capability query advertised DMA-BUF, but the weakly-linked entry point
    // did not resolve (ROCr runtime too old to export it). Disable the feature
    // cleanly rather than leaving an inconsistent gate.
    INFO(NCCL_INIT, "ROCr runtime does not export hsa_amd_portable_export_dmabuf; disabling DMA-BUF.");
    goto error;
  }
  else {
    // Arm the DMA-BUF feature gate with the resolved HSA symbol.
    pfn_hsa_amd_portable_export_dmabuf = hsa_amd_portable_export_dmabuf;
  }

  //check OS kernel support
  if(!rcclParamForceEnableDMABUF()) {
    struct utsname utsname;
    FILE *fp = NULL;
    char kernel_opt1[28] = "CONFIG_DMABUF_MOVE_NOTIFY=y";
    char kernel_opt2[20] = "CONFIG_PCI_P2PDMA=y";
    char kernel_conf_file[128];
    char buf[256];
    int found_opt1 = 0;
    int found_opt2 = 0;

    //check for kernel name exists
    if (uname(&utsname) == -1) INFO(NCCL_INIT,"Could not get kernel name");
    //format and store the kernel conf file location
    const char* possiblePaths[] = {
      "/proc/config.gz",
      "/boot/config-%s",
      "/usr/src/linux-%s/.config",
      "/usr/src/linux/.config",
      "/usr/lib/modules/%s/config",
      "/usr/lib/ostree-boot/config-%s",
      "/usr/lib/kernel/config-%s",
      "/usr/src/linux-headers-%s/.config",
      "/lib/modules/%s/build/.config",
    };

    // Check if zcat is available in the system
    int has_zcat = (system("which zcat > /dev/null 2>&1") == 0);

    for (const auto& path : possiblePaths) {
      // Reset flags for each file
      found_opt1 = 0;
      found_opt2 = 0;

      // Special handling for /proc/config.gz
      snprintf(kernel_conf_file, sizeof(kernel_conf_file), path, utsname.release);

      if (strstr(path, "/proc/config.gz") != NULL) {
        // Skip if zcat is unavailable or /proc/config.gz does not exist.
        // popen() succeeds even when the file is missing, producing an empty
        // stream that falsely triggers the "not found" error path.
        if (!has_zcat || access("/proc/config.gz", R_OK) != 0) {
          INFO(NCCL_INIT, "Skipping %s (zcat %s, file %s)", kernel_conf_file,
               has_zcat ? "available" : "unavailable",
               access("/proc/config.gz", R_OK) == 0 ? "exists" : "not found");
          continue;
        }
        fp = popen("zcat /proc/config.gz 2>/dev/null", "r");
      } else {
        fp = fopen(kernel_conf_file, "r");
      }

      if (fp != NULL){
        //look for kernel_opt1 and kernel_opt2 in the conf file and check
        while (fgets(buf, sizeof(buf), fp) != NULL) {
          if (strstr(buf, kernel_opt1) != NULL) {
            found_opt1 = 1;
            INFO(NCCL_INIT,"%s in %s", kernel_opt1, kernel_conf_file);
          }
          if (strstr(buf, kernel_opt2) != NULL) {
            found_opt2 = 1;
            INFO(NCCL_INIT,"%s in %s", kernel_opt2, kernel_conf_file);
          }
        }

        // Close file handle
        if (strstr(path, "/proc/config.gz") != NULL) {
          pclose(fp);
        } else {
          fclose(fp);
        }

        // Check if both options were found
        if (!found_opt1 || !found_opt2) {
          dmaBufSupport = 0;
          INFO(NCCL_INIT, "CONFIG_DMABUF_MOVE_NOTIFY and CONFIG_PCI_P2PDMA should be set for DMA_BUF in %s", kernel_conf_file);
          INFO(NCCL_INIT, "DMA_BUF_SUPPORT Failed due to OS kernel support");
        }

        if(dmaBufSupport) INFO(NCCL_INIT, "DMA_BUF Support Enabled");
        else goto error;
        break;
      }
    }
    if(fp == NULL) {
      // Fallback: check /proc/kallsyms for DMA-BUF and P2PDMA kernel symbols.
      // Works inside Docker containers where /boot/config-* is unavailable.
      INFO(NCCL_INIT, "Could not open kernel conf file, trying /proc/kallsyms fallback");
      FILE *kallsyms = fopen("/proc/kallsyms", "r");
      if (kallsyms) {
        while (fgets(buf, sizeof(buf), kallsyms) != NULL) {
          if (!found_opt1 && strstr(buf, "dma_buf_move_notify") != NULL)
            found_opt1 = 1;
          if (!found_opt2 && strstr(buf, "pci_p2pdma") != NULL)
            found_opt2 = 1;
          if (found_opt1 && found_opt2) break;
        }
        fclose(kallsyms);
        if (found_opt1 && found_opt2) {
          INFO(NCCL_INIT, "DMA_BUF Support Enabled via /proc/kallsyms (dma_buf_move_notify + pci_p2pdma)");
        } else {
          dmaBufSupport = 0;
          INFO(NCCL_INIT, "DMA_BUF_SUPPORT Failed: missing kernel symbols in /proc/kallsyms");
          goto error;
        }
      } else {
        dmaBufSupport = 0;
        INFO(NCCL_INIT, "Could not open /proc/kallsyms");
      }
    }
  }
  /*
   * Required to initialize the ROCr Driver.
   * Multiple calls of hsa_init() will return immediately
   * without making any relevant change
   */
  hsa_init();

  initResult = ncclSuccess;
  return;

error:
  initResult = ncclSystemError;
}

ncclResult_t rocmLibraryInit() {
  pthread_once(&initOnceControl, initOnceFunc);
  return initResult;
}
