nvrtc_chunk_fusion.h

Go to the documentation of this file.

#pragma once
 
#include "fused/chunk_fusion/chunk_fusion.h"   // ChunkCoderKind, chunk geometry
#include "backend/types.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
 
namespace fz {
namespace fused {
 
struct ChunkFusionSpec {
    std::string              quant_op    = "QuantInplaceZigzag";
    std::vector<std::string> transforms  = {"DiffNegabinary", "Bitshuffle32"};
    std::string              coder       = "RZECoder";
    int                       chunk_bytes = 16384;
};
 
inline const char* chunkCoderOpName(ChunkCoderKind k) {
    return (k == ChunkCoderKind::RRE) ? "RRECoder" : "RZECoder";
}
 
bool nvrtcChunkFusionAvailable();
 
std::string generateChunkFusionSource(const ChunkFusionSpec& spec);
 
void launchNvrtcChunkFusedEncode(
    const ChunkFusionSpec& spec, const float* d_in, size_t n, const uint8_t* d_params,
    uint8_t* d_scratch, uint32_t* d_sizes, unsigned nc, fz::stream_t stream,
    uint32_t* d_side_idxs = nullptr, float* d_side_vals = nullptr,
    uint32_t* d_side_count = nullptr, uint32_t side_max = 0);
 
size_t launchGenericChunkFusion(
    const ChunkFusionSpec& spec, const float* d_in, size_t n,
    const uint8_t* host_params, size_t params_bytes,
    uint8_t* d_out, MemoryPool* pool, fz::stream_t stream,
    uint32_t* d_side_idxs = nullptr, float* d_side_vals = nullptr,
    uint32_t side_max = 0, uint32_t* out_side_count = nullptr);
 
} // namespace fused
} // namespace fz