#pragma once

#include "common.h"
#include "Tensor.h"
#include "kernels/zgemm/zgemm.h"

namespace nunchaku::utils {

void set_cuda_stack_limit(int64_t newval) {
    size_t val = 0;
    checkCUDA(hipDeviceSetLimit(hipLimitStackSize, (size_t)newval));
    checkCUDA(hipDeviceGetLimit(&val, hipLimitStackSize));
    spdlog::debug("Stack={}", val);
}

void disable_memory_auto_release() {
    int device;
    checkCUDA(hipGetDevice(&device));
    hipMemPool_t mempool;
    checkCUDA(hipDeviceGetDefaultMemPool(&mempool, device));
    uint64_t threshold = UINT64_MAX;
    checkCUDA(hipMemPoolSetAttribute(mempool, hipMemPoolAttrReleaseThreshold, &threshold));
}

void trim_memory() {
    int device;
    checkCUDA(hipGetDevice(&device));
    hipMemPool_t mempool;
    checkCUDA(hipDeviceGetDefaultMemPool(&mempool, device));
    size_t bytesToKeep = 0;
    checkCUDA(hipMemPoolTrimTo(mempool, bytesToKeep));
}

void set_faster_i2f_mode(std::string mode) {
    spdlog::info("Set fasteri2f mode to {}", mode);
    kernels::set_faster_i2f_mode(mode);
}

}; // namespace nunchaku::utils
