#include <cuda_runtime.h>
#include <cstdint>
#include <cstring>
#include <memory>
#include <mutex>
#include <new>
#include <string>
#include "nvcomp/bitcomp.hpp"
#include "nvcomp/zstd.hpp"
#include "nvcomp/nvcompManager.hpp"
using namespace nvcomp;
namespace {
constexpr int FERRO_SHIM_OK = 0;
constexpr int FERRO_SHIM_NULL_HANDLE = 0x0001'0000;
constexpr int FERRO_SHIM_BAD_ARG = 0x0002'0000;
constexpr int FERRO_SHIM_CUDA_ERROR = 0x0003'0000;
constexpr int FERRO_SHIM_NVCOMP_ERR = 0x0004'0000;
constexpr int FERRO_SHIM_BAD_ALLOC = 0x0005'0000;
constexpr int FERRO_SHIM_EXCEPTION = 0x0006'0000;
thread_local std::string g_last_error_msg;
void set_err(const std::string& msg)
{
g_last_error_msg = msg;
}
void clear_err()
{
g_last_error_msg.clear();
}
struct ManagerHandle {
std::shared_ptr<nvcompManagerBase> manager;
cudaStream_t stream = nullptr;
size_t chunk_size = 0;
std::unique_ptr<CompressionConfig> cached_comp_cfg;
size_t cached_comp_cfg_size = 0;
};
}
extern "C" {
int ferro_nvcomp_hlif_create_bitcomp(
size_t chunk_size,
int algorithm,
nvcompType_t data_type,
cudaStream_t user_stream,
void** out_handle)
{
clear_err();
if (out_handle == nullptr) {
set_err("ferro_nvcomp_hlif_create_bitcomp: out_handle is null");
return FERRO_SHIM_BAD_ARG;
}
*out_handle = nullptr;
if (chunk_size == 0 || chunk_size > nvcompBitcompCompressionMaxAllowedChunkSize) {
set_err("ferro_nvcomp_hlif_create_bitcomp: chunk_size out of range");
return FERRO_SHIM_BAD_ARG;
}
if (algorithm != 0 && algorithm != 1) {
set_err("ferro_nvcomp_hlif_create_bitcomp: algorithm must be 0 or 1");
return FERRO_SHIM_BAD_ARG;
}
try {
nvcompBatchedBitcompCompressOpts_t comp_opts{algorithm, data_type, {0}};
auto mgr = std::make_shared<BitcompManager>(
chunk_size,
comp_opts,
nvcompBatchedBitcompDecompressDefaultOpts,
user_stream,
NoComputeNoVerify);
auto handle = new (std::nothrow) ManagerHandle();
if (handle == nullptr) {
set_err("ferro_nvcomp_hlif_create_bitcomp: bad_alloc on ManagerHandle");
return FERRO_SHIM_BAD_ALLOC;
}
handle->manager = std::move(mgr);
handle->stream = user_stream;
handle->chunk_size = chunk_size;
*out_handle = handle;
return FERRO_SHIM_OK;
} catch (const std::bad_alloc& e) {
set_err(std::string("BitcompManager: bad_alloc: ") + e.what());
return FERRO_SHIM_BAD_ALLOC;
} catch (const std::exception& e) {
set_err(std::string("BitcompManager: exception: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("BitcompManager: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
int ferro_nvcomp_hlif_create_zstd(
size_t chunk_size,
cudaStream_t user_stream,
void** out_handle)
{
clear_err();
if (out_handle == nullptr) {
set_err("ferro_nvcomp_hlif_create_zstd: out_handle is null");
return FERRO_SHIM_BAD_ARG;
}
*out_handle = nullptr;
if (chunk_size == 0 || chunk_size > (1 << 24)) {
set_err("ferro_nvcomp_hlif_create_zstd: chunk_size out of range");
return FERRO_SHIM_BAD_ARG;
}
try {
auto mgr = std::make_shared<ZstdManager>(
chunk_size,
nvcompBatchedZstdCompressDefaultOpts,
nvcompBatchedZstdDecompressDefaultOpts,
user_stream,
NoComputeNoVerify);
auto handle = new (std::nothrow) ManagerHandle();
if (handle == nullptr) {
set_err("ferro_nvcomp_hlif_create_zstd: bad_alloc on ManagerHandle");
return FERRO_SHIM_BAD_ALLOC;
}
handle->manager = std::move(mgr);
handle->stream = user_stream;
handle->chunk_size = chunk_size;
*out_handle = handle;
return FERRO_SHIM_OK;
} catch (const std::bad_alloc& e) {
set_err(std::string("ZstdManager: bad_alloc: ") + e.what());
return FERRO_SHIM_BAD_ALLOC;
} catch (const std::exception& e) {
set_err(std::string("ZstdManager: exception: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("ZstdManager: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
void ferro_nvcomp_hlif_destroy(void* handle)
{
if (handle == nullptr) {
return;
}
auto* h = static_cast<ManagerHandle*>(handle);
h->cached_comp_cfg.reset();
h->manager.reset();
delete h;
}
int ferro_nvcomp_hlif_max_compressed_size(
void* handle,
size_t uncomp_bytes,
size_t* out_max_bytes)
{
clear_err();
if (handle == nullptr || out_max_bytes == nullptr) {
set_err("ferro_nvcomp_hlif_max_compressed_size: null handle or out_max_bytes");
return FERRO_SHIM_NULL_HANDLE;
}
auto* h = static_cast<ManagerHandle*>(handle);
try {
auto cfg = std::make_unique<CompressionConfig>(
h->manager->configure_compression(uncomp_bytes));
*out_max_bytes = cfg->max_compressed_buffer_size;
h->cached_comp_cfg = std::move(cfg);
h->cached_comp_cfg_size = uncomp_bytes;
return FERRO_SHIM_OK;
} catch (const std::exception& e) {
set_err(std::string("configure_compression: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("configure_compression: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
int ferro_nvcomp_hlif_compress(
void* handle,
const uint8_t* d_in,
size_t uncomp_bytes,
uint8_t* d_out,
size_t* out_comp_bytes)
{
clear_err();
if (handle == nullptr || out_comp_bytes == nullptr) {
set_err("ferro_nvcomp_hlif_compress: null handle or out_comp_bytes");
return FERRO_SHIM_NULL_HANDLE;
}
if (d_in == nullptr || d_out == nullptr) {
set_err("ferro_nvcomp_hlif_compress: null device pointer");
return FERRO_SHIM_BAD_ARG;
}
auto* h = static_cast<ManagerHandle*>(handle);
try {
if (!h->cached_comp_cfg || h->cached_comp_cfg_size != uncomp_bytes) {
h->cached_comp_cfg = std::make_unique<CompressionConfig>(
h->manager->configure_compression(uncomp_bytes));
h->cached_comp_cfg_size = uncomp_bytes;
}
h->manager->compress(d_in, d_out, *h->cached_comp_cfg);
*out_comp_bytes = h->manager->get_compressed_output_size(d_out);
return FERRO_SHIM_OK;
} catch (const std::exception& e) {
set_err(std::string("compress: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("compress: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
int ferro_nvcomp_hlif_decompress(
void* handle,
const uint8_t* d_comp,
uint8_t* d_out,
size_t* out_decomp_bytes)
{
clear_err();
if (handle == nullptr) {
set_err("ferro_nvcomp_hlif_decompress: null handle");
return FERRO_SHIM_NULL_HANDLE;
}
if (d_comp == nullptr || d_out == nullptr) {
set_err("ferro_nvcomp_hlif_decompress: null device pointer");
return FERRO_SHIM_BAD_ARG;
}
auto* h = static_cast<ManagerHandle*>(handle);
try {
auto decomp_cfg = h->manager->configure_decompression(d_comp);
h->manager->decompress(d_out, d_comp, decomp_cfg);
if (out_decomp_bytes != nullptr) {
*out_decomp_bytes = decomp_cfg.decomp_data_size;
}
return FERRO_SHIM_OK;
} catch (const std::exception& e) {
set_err(std::string("decompress: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("decompress: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
int ferro_nvcomp_hlif_get_decompressed_output_size(
void* handle,
const uint8_t* d_comp,
size_t* out_decomp_bytes)
{
clear_err();
if (handle == nullptr || d_comp == nullptr || out_decomp_bytes == nullptr) {
set_err("ferro_nvcomp_hlif_get_decompressed_output_size: null arg");
return FERRO_SHIM_NULL_HANDLE;
}
auto* h = static_cast<ManagerHandle*>(handle);
try {
*out_decomp_bytes = h->manager->get_decompressed_output_size(d_comp);
return FERRO_SHIM_OK;
} catch (const std::exception& e) {
set_err(std::string("get_decompressed_output_size: ") + e.what());
return FERRO_SHIM_EXCEPTION;
} catch (...) {
set_err("get_decompressed_output_size: unknown exception");
return FERRO_SHIM_EXCEPTION;
}
}
size_t ferro_nvcomp_hlif_last_error_message(char* buf, size_t buf_capacity)
{
const std::string& msg = g_last_error_msg;
if (buf == nullptr || buf_capacity == 0) {
return msg.size();
}
const size_t copy_len = msg.size() < (buf_capacity - 1) ? msg.size() : (buf_capacity - 1);
std::memcpy(buf, msg.data(), copy_len);
buf[copy_len] = '\0';
return msg.size();
}
}