THSTorch.cpp

// Copyright (c) .NET Foundation and Contributors.  All Rights Reserved.  See LICENSE in the project root for license information.
#include "THSTorch.h"

#include "torch/torch.h"
#include "torch/cuda.h"
#include <c10/core/CachingDeviceAllocator.h>

#if defined(USE_CUDA)
#include <c10/cuda/CUDACachingAllocator.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAFunctions.h>
#endif

void THSTorch_manual_seed(const int64_t seed)
{
    torch::manual_seed(seed);
}

Generator THSGenerator_manual_seed(const int64_t seed)
{
    auto gen = at::globalContext().defaultGenerator(at::DeviceType::CPU);
    gen.set_current_seed(seed);
    return new at::Generator(gen.getIntrusivePtr());
}

void THSCuda_manual_seed(const int64_t seed)
{
    CATCH(torch::cuda::manual_seed(seed);)
}

void THSCuda_manual_seed_all(const int64_t seed)
{
    CATCH(torch::cuda::manual_seed_all(seed);)
}

bool THSBackend_cublas_get_allow_tf32()
{
    auto result = false;
    CATCH(result = at::globalContext().allowTF32CuBLAS(););
    return result;
}

void THSBackend_cublas_set_allow_tf32(const bool flag)
{
    CATCH(at::globalContext().setAllowTF32CuBLAS(flag););
}

bool THSBackend_cudnn_get_allow_tf32()
{
    auto result = false;
    CATCH(result = at::globalContext().allowTF32CuDNN(););
    return result;
}

void THSBackend_cudnn_set_allow_tf32(const bool flag)
{
    CATCH(at::globalContext().setAllowTF32CuDNN(flag););
}

bool THSBackend_cuda_get_allow_fp16_reduced_precision_reduction()
{
    auto result = false;
    CATCH(result = at::globalContext().allowFP16ReductionCuBLAS() == at::CuBLASReductionOption::AllowReducedPrecisionWithSplitK;);
    return result;
}

void THSBackend_cuda_set_allow_fp16_reduced_precision_reduction(const bool flag)
{
    CATCH(at::globalContext().setAllowFP16ReductionCuBLAS(flag, true););
}

bool THSBackend_cuda_get_enable_flash_sdp()
{
    auto result = false;
    CATCH(result = at::globalContext().userEnabledFlashSDP(););
    return result;
}

void THSBackend_cuda_set_enable_flash_sdp(const bool flag)
{
    CATCH(at::globalContext().setSDPUseFlash(flag););
}

bool THSBackend_cuda_get_enable_math_sdp()
{
    auto result = false;
    CATCH(result = at::globalContext().userEnabledMathSDP(););
    return result;
}

void THSBackend_cuda_set_enable_math_sdp(const bool flag)
{
    CATCH(at::globalContext().setSDPUseMath(flag););
}

void THSGenerator_gen_manual_seed(const Generator generator, const int64_t seed)
{
    generator->set_current_seed(seed);
}

Generator THSGenerator_default_generator()
{
    auto gen = at::globalContext().defaultGenerator(at::DeviceType::CPU);
    return new at::Generator(gen.getIntrusivePtr());
}

int64_t THSGenerator_initial_seed(const Generator gen)
{
    return gen->current_seed();
}

Tensor THSGenerator_get_rng_state(const Generator gen)
{
    CATCH_TENSOR(gen->get_state());
}

void  THSGenerator_set_rng_state(const Generator gen, const Tensor tensor)
{
    gen->set_state(*tensor);
}


Generator THSGenerator_new(uint64_t seed, int64_t device, int64_t index)
{
    // TODO: Support creation of GPU RNGs. 'device' and 'index' are in the
    //       function signature in preparation thereof.
    return new at::Generator(at::detail::createCPUGenerator(seed));
}

void THSGenerator_dispose(const Generator generator)
{
    delete generator;
}

int THSTorchCuda_is_available()
{
    return torch::cuda::is_available();
}

int THSTorchCuda_cudnn_is_available()
{
    return torch::cuda::cudnn_is_available();
}

int THSTorchCuda_device_count()
{
    return (int)torch::cuda::device_count();
}

void THSTorchCuda_synchronize(const int64_t device_index)
{
    CATCH(torch::cuda::synchronize(device_index);)
}

#if defined(USE_CUDA)

// Helper to safely resolve CUDA device index, defaulting to device 0 when
// no device context has been established yet (c10::cuda::current_device()
// can throw "Invalid device argument" in that case).
static c10::DeviceIndex resolve_cuda_device_index(int64_t device_index)
{
    if (device_index >= 0)
        return static_cast<c10::DeviceIndex>(device_index);
    c10::DeviceIndex current = -1;
    auto err = c10::cuda::GetDevice(&current);
    if (err != cudaSuccess || current < 0) {
        // No CUDA device context has been established yet; default to device 0.
        cudaGetLastError(); // clear any error
        return static_cast<c10::DeviceIndex>(0);
    }
    return current;
}

void THSTorchCuda_empty_cache()
{
    CATCH(c10::cuda::CUDACachingAllocator::emptyCache();)
}

size_t THSTorchCuda_memory_allocated(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        res = c10::cuda::CUDACachingAllocator::currentMemoryAllocated(device);
    )
    return res;
}

size_t THSTorchCuda_max_memory_allocated(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        res = c10::cuda::CUDACachingAllocator::maxMemoryAllocated(device);
    )
    return res;
}

void THSTorchCuda_reset_peak_memory_stats(const int64_t device_index)
{
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        c10::cuda::CUDACachingAllocator::resetPeakStats(device);
    )
}

size_t THSTorchCuda_memory_reserved(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        res = c10::cuda::CUDACachingAllocator::currentMemoryReserved(device);
    )
    return res;
}

size_t THSTorchCuda_max_memory_reserved(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        res = c10::cuda::CUDACachingAllocator::maxMemoryReserved(device);
    )
    return res;
}

void THSTorchCuda_mem_get_info(const int64_t device_index, size_t* free, size_t* total)
{
    CATCH(
        auto device = resolve_cuda_device_index(device_index);
        c10::cuda::CUDAGuard guard(device);
        C10_CUDA_CHECK(cudaMemGetInfo(free, total));
    )
}

void THSTorchCuda_set_device(const int64_t device_index)
{
    CATCH(c10::cuda::set_device(static_cast<c10::DeviceIndex>(device_index));)
}

int64_t THSTorchCuda_current_device()
{
    int64_t res = -1;
    CATCH(res = static_cast<int64_t>(c10::cuda::current_device());)
    return res;
}

#else

// Helper to resolve device index using torch::cuda APIs available without USE_CUDA
static c10::DeviceIndex resolve_cuda_device(int64_t device_index)
{
    if (device_index < 0) {
        // Use device 0 as default when CUDA headers aren't available
        return 0;
    }
    return static_cast<c10::DeviceIndex>(device_index);
}

void THSTorchCuda_empty_cache()
{
    CATCH(
        if (torch::cuda::is_available()) {
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            allocator->emptyCache();
        }
    )
}

size_t THSTorchCuda_memory_allocated(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            auto stats = allocator->getDeviceStats(device);
            res = static_cast<size_t>(stats.allocated_bytes[static_cast<size_t>(c10::CachingAllocator::StatType::AGGREGATE)].current);
        }
    )
    return res;
}

size_t THSTorchCuda_max_memory_allocated(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            auto stats = allocator->getDeviceStats(device);
            res = static_cast<size_t>(stats.allocated_bytes[static_cast<size_t>(c10::CachingAllocator::StatType::AGGREGATE)].peak);
        }
    )
    return res;
}

void THSTorchCuda_reset_peak_memory_stats(const int64_t device_index)
{
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            allocator->resetPeakStats(device);
        }
    )
}

size_t THSTorchCuda_memory_reserved(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            auto stats = allocator->getDeviceStats(device);
            res = static_cast<size_t>(stats.reserved_bytes[static_cast<size_t>(c10::CachingAllocator::StatType::AGGREGATE)].current);
        }
    )
    return res;
}

size_t THSTorchCuda_max_memory_reserved(const int64_t device_index)
{
    size_t res = 0;
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            auto stats = allocator->getDeviceStats(device);
            res = static_cast<size_t>(stats.reserved_bytes[static_cast<size_t>(c10::CachingAllocator::StatType::AGGREGATE)].peak);
        }
    )
    return res;
}

void THSTorchCuda_mem_get_info(const int64_t device_index, size_t* free, size_t* total)
{
    *free = 0;
    *total = 0;
    CATCH(
        if (torch::cuda::is_available()) {
            auto device = resolve_cuda_device(device_index);
            auto* allocator = c10::getDeviceAllocator(c10::DeviceType::CUDA);
            auto info = allocator->getMemoryInfo(device);
            *free = info.first;
            *total = info.second;
        }
    )
}

void THSTorchCuda_set_device(const int64_t device_index)
{
}

int64_t THSTorchCuda_current_device()
{
    return -1;
}

#endif


const char * THSTorch_get_and_reset_last_err()
{
    char *tmp = torch_last_err;
    torch_last_err = nullptr;
    return tmp;
}

int THSTorch_get_num_threads()
{
    CATCH_RETURN_RES(int, -1, res = torch::get_num_threads());
}

void THSTorch_set_num_threads(const int threads)
{
    torch::set_num_threads(threads);
}

int THSTorch_get_num_interop_threads()
{
    CATCH_RETURN_RES(int, -1, res = torch::get_num_interop_threads());
}

void THSTorch_set_num_interop_threads(const int threads)
{
    torch::set_num_interop_threads(threads);
}

int THSTorch_can_cast(const int type1, const int type2)
{
    CATCH_RETURN_RES(int, -1, res = (int)torch::can_cast((c10::ScalarType)type1, (c10::ScalarType)type2));
}

int THSTorch_promote_types(const int type1, const int type2)
{
    CATCH_RETURN_RES(int, -1, res = (int)torch::promote_types((c10::ScalarType)type1, (c10::ScalarType)type2));
}


Scalar THSTorch_int8_to_scalar(int8_t value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_uint8_to_scalar(uint8_t value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_int16_to_scalar(int16_t value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_int32_to_scalar(int32_t value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_int64_to_scalar(int64_t value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_float32_to_scalar(float value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_float64_to_scalar(double value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_float16_to_scalar(float value)
{
    return new torch::Scalar((c10::Half)value);
}

Scalar THSTorch_bfloat16_to_scalar(float value)
{
    return new torch::Scalar((c10::BFloat16)value);
}

Scalar THSTorch_bool_to_scalar(bool value)
{
    return new torch::Scalar(value);
}

Scalar THSTorch_complex32_to_scalar(float real, float imaginary)
{
    return new torch::Scalar(c10::complex<float>(real, imaginary));
}

Scalar THSTorch_complex64_to_scalar(double real, double imaginary)
{
    return new torch::Scalar(c10::complex<double>(real, imaginary));
}

int8_t THSTorch_scalar_to_int8(Scalar value)
{
    return value->toChar();
}

uint8_t THSTorch_scalar_to_uint8(Scalar value)
{
    return value->toByte();
}

int16_t THSTorch_scalar_to_int16(Scalar value)
{
    return value->toShort();
}

int32_t THSTorch_scalar_to_int32(Scalar value)
{
    return value->toInt();
}

int64_t THSTorch_scalar_to_int64(Scalar value)
{
    return value->toLong();
}

float THSTorch_scalar_to_float32(Scalar value)
{
    return value->toFloat();
}

double THSTorch_scalar_to_float64(Scalar value)
{
    return value->toDouble();
}

void THSTorch_scalar_to_bfloat16(Scalar value, unsigned short* res)
{
    *res = value->toBFloat16().x;
}

void THSTorch_scalar_to_float16(Scalar value, unsigned short *res)
{
    *res = value->toHalf().x;
}

void THSTorch_scalar_to_complex32(Scalar value, float* real, float* imaginary)
{
    auto result = value->toComplexFloat();
    *real = result.real();
    *imaginary = result.imag();
}

void THSTorch_scalar_to_complex64(Scalar value, double* real, double* imaginary)
{
    auto result = value->toComplexDouble();
    *real = result.real();
    *imaginary = result.imag();
}

bool THSTorch_scalar_to_bool(Scalar value)
{
    return value->toBool();
}

int8_t THSTorch_scalar_type(Scalar value)
{
    return (int8_t)value->type();
}

void THSTorch_dispose_scalar(Scalar scalar)
{
    delete scalar;
}

double THSSpecial_erf_scalar(const double x)
{
    return erf(x);
}

double THSSpecial_erfc_scalar(const double x)
{
    return erfc(x);
}