nvrtc_warp_fusion.h

Go to the documentation of this file.

#pragma once
 
#include "backend/types.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
 
namespace fz {
class MemoryPool;
 
namespace fused {
 
struct WarpFusionSpec {
    std::string              predictor      = "Lorenzo1DPredictor";
    std::vector<std::string> transforms;                        // optional, in execution order
    std::string              coder          = "AdaptiveBitpackCoder";   // swappable sink
    int                      elems_per_lane = 1;
 
    std::string              predictor_ti;
    std::string              coder_ti;
};
 
std::string generateWarpFusionSource(const WarpFusionSpec& spec);
 
size_t launchNvrtcWarpFused(
    const WarpFusionSpec& spec, const float* d_in, size_t n_ab,
    const uint8_t* pred_params, size_t params_bytes,
    uint8_t* d_out, MemoryPool* pool, fz::stream_t stream,
    std::string* execution_path = nullptr);
 
size_t launchNvrtcWarpInverseFused(
    const WarpFusionSpec& spec, const uint8_t* d_archive, size_t archive_bytes,
    size_t n_elems, size_t n_out, float ebx2,
    const uint8_t* pred_params, size_t params_bytes,
    float* d_out, MemoryPool* pool, fz::stream_t stream);
 
std::string generateWarpInverseSource(const WarpFusionSpec& spec);
 
} // namespace fused
} // namespace fz