//
// Created by xuwei on 2024/2/8.
//

#include <sys/types.h>
#include <iostream>
#include <torch/extension.h>

#include "acl_base.h"
#include "acl_rt.h"
#include "PluggableAllocator.h"
#include "Decorator.h"

extern "C" {
void *my_malloc(size_t size, int device, aclrtStream stream) {
    void *ptr;
    ptr = PluggableAllocator::getInstance().malloc(device, size, stream);
    std::cout << "alloc ptr = " << ptr << ", size = " << size << std::endl;
    return ptr;
}

void my_free(void *ptr, size_t size, int device, aclrtStream stream) {
    PluggableAllocator::getInstance().free(ptr);
    std::cout << "pluggable free success" << std::endl;
}

void my_init(int device_count) {
    PluggableAllocator::getInstance().init(device_count);
    std::cout << "device: " << device_count << "init success" << std::endl;
}

void my_empty_cache(bool check_error) {
    PluggableAllocator::getInstance().emptyCache(true);
    std::cout << "pluggable empty cache success" << std::endl;
}

void my_memory_fraction(double fraction, int device) {
    PluggableAllocator::getInstance().setMemoryFraction(fraction, device);
    std::cout << "using pluggable memory_fraction success" << std::endl;
}

DeviceStats my_get_device_stats(int device) {
    std::cout << "using pluggable get_device_stats success" << std::endl;
    return PluggableAllocator::getInstance().getDeviceStats(device);
}

//void *my_base_alloc(void *ptr, size_t *size) {
//    return PluggableAllocator::getInstance().getBaseAllocation(ptr, size);
//}

// TODO 接口有问题
//void my_record_stream(const c10::DataPtr &ptr, c10_npu::NPUStream stream) {
//    Allocator::NpuCachingCustomAllocator::getInstance().recordStream(ptr, stream);
//}

// TODO 接口有问题
//void my_erase_stream(const c10::DataPtr &ptr, c10_npu::NPUStream stream) {
//    Allocator::NpuCachingCustomAllocator::getInstance().eraseStream(ptr, stream);
//}
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("memory_recorder_start", &Decorator::memory_recorder_start, "start mark the life cycle of a tensor in forward");
    m.def("memory_recorder_end", &Decorator::memory_recorder_end, "end mark the life cycle of a tensor in forward");
    m.def("malloc_recorder_start", &Decorator::malloc_recorder_start, "start mark the life cycle of a tensor in step");
    m.def("malloc_recorder_end", &Decorator::malloc_recorder_end, "end mark the life cycle of a tensor in step");
    m.def("precise_match_start", &Decorator::precise_match_start, "start mark the life cycle of a tensor in optimizer init stage");
    m.def("precise_match_end", &Decorator::precise_match_end, "end mark the life cycle of a tensor in optimizer init stage");
}