rtc.h

Go to the documentation of this file.

#pragma once

#include <chrono>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

#include <cuda.h>
#include <driver_types.h> // cuda driver types

namespace tc {

extern std::mutex nvrtc_mutex;
using Duration = std::chrono::high_resolution_clock::duration;

//
// Basic interface to expose NVRTC JIT compilation and module
// loading/unloading + API kernel launches.
//
class CudaRTCFunction {
  CudaRTCFunction();

 public:
  ~CudaRTCFunction();

  static std::shared_ptr<CudaRTCFunction> Compile(
      const std::string& name,
      const std::string& source);

  // if profile is set it returns the kernel runtime
  Duration Launch(
      const std::array<size_t, 3>& grid,
      const std::array<size_t, 3>& block,
      unsigned int shared_mem,
      cudaStream_t stream,
      // by copy because we take an address to element when calling the kernel
      // TODO: check the overhead of double indirection on kernel calls, this
      // does not look ideal for low-latency
      std::vector<int> params,
      std::vector<void*> outputs,
      std::vector<const void*> inputs,
      bool profile = false) const;

  void clear();

 private:
  mutable std::unordered_map<size_t, CUmodule> perGpuModule_;
  mutable std::unordered_map<size_t, CUfunction> perGpuKernel_;
  std::string specializedName;
  std::vector<char> nvrtc_ptx;
  bool cleared_;
};

} // namespace tc