/include/dgl/runtime/tensordispatch.h

/*!
 *  Copyright (c) 2020-2022 by Contributors
 * \file array/tensordispatch.h
 * \brief This file defines the dispatcher of tensor operators to framework-specific
 *  implementations.
 *
 *  The dispatcher consists of a TensorDispatcher singleton in DGL C library and
 *  one separately-built shared library per supported backend.
 *
 *  Those shared libraries contain wrappers of the framework-specific operators.
 *  The wrappers are defined with extern "C", meaning that the C++ compiler will
 *  not do name mangling for those functions so that DGL can conveniently locate
 *  them using dlsym(3) (or GetProcAddress in Windows).
 *
 *  The TensorDispatcher singleton maintains a mapping from an array operator to
 *  the address of the corresponding symbol in the shared library.  During
 *  initialization, the TensorDispatcher checks which backend DGL is using.
 *  It then locates and opens the corresponding shared library using dlopen(3) (or
 *  LoadLibrary in Windows), and populates the said mapping above with dlsym(3)
 *  (or GetProcAddress in Windows).
 *
 *  A tensor operator in TensorDispatcher first checks whether the corresponding symbol
 *  address is found in the mapping.  If so, it calls the function located at the
 *  symbol address instead, allocate/free pieces of memory on CPU/GPU.
 *  If not, it falls back to DeviceAPI::AllocWorkspace/FreeWorkspace.
 */

#ifndef DGL_RUNTIME_TENSORDISPATCH_H_
#define DGL_RUNTIME_TENSORDISPATCH_H_

#include <stddef.h>
#include <tensoradapter.h>
#if defined(WIN32) || defined(_WIN32)
#include <windows.h>
#endif  // WIN32
#ifdef DGL_USE_CUDA
#include <cuda_runtime.h>
#endif  // DGL_USE_CUDA
#include "ndarray.h"

/*! \brief Casts a pointer \c entry to a function pointer with signature of \c func */
#define FUNCCAST(func, entry)   (*reinterpret_cast<decltype(&(func))>(entry))

namespace dgl {
namespace runtime {

/*!
 * \brief Dispatcher that delegates the function calls to framework-specific C++ APIs.
 *
 * This class is not thread-safe.
 */
class TensorDispatcher {
 public:
  /*! \brief Get the singleton instance. */
  static TensorDispatcher* Global() {
    static TensorDispatcher inst;
    return &inst;
  }

  /*! \brief Whether an adapter library is available */
  inline bool IsAvailable() {
    return available_;
  }

  /*! \brief Load symbols from the given tensor adapter library path */
  bool Load(const char *path_cstr);

  /*!
   * \brief Allocate a piece of CPU memory via
   * PyTorch's CPUAllocator.
   * Used in CPUDeviceAPI::AllocWorkspace().
   *
   * \param nbytes The size to be allocated.
   * \return Pointer to the allocated memory.
   */
  inline void* CPUAllocWorkspace(size_t nbytes) {
    auto entry = entrypoints_[Op::kCPURawAlloc];
    return FUNCCAST(tensoradapter::CPURawAlloc, entry)(nbytes);
  }

  /*!
   * \brief Free the CPU memory.
   * Used in CPUDeviceAPI::FreeWorkspace().
   *
   * \param ptr Pointer to the memory to be freed.
   */
  inline void CPUFreeWorkspace(void* ptr) {
    auto entry = entrypoints_[Op::kCPURawDelete];
    FUNCCAST(tensoradapter::CPURawDelete, entry)(ptr);
  }

#ifdef DGL_USE_CUDA
  /*!
   * \brief Allocate a piece of GPU memory via
   * PyTorch's THCCachingAllocator.
   * Used in CUDADeviceAPI::AllocWorkspace().
   * 
   * \note THCCachingAllocator specify the device to allocate on
   * via cudaGetDevice(). Make sure to call cudaSetDevice()
   * before invoking this function.
   *
   * \param nbytes The size to be allocated.
   * \param stream The stream to be allocated on.
   * \return Pointer to the allocated memory.
   */
  inline void* CUDAAllocWorkspace(size_t nbytes, cudaStream_t stream) {
    auto entry = entrypoints_[Op::kCUDARawAlloc];
    return FUNCCAST(tensoradapter::CUDARawAlloc, entry)(nbytes, stream);
  }

  /*!
   * \brief Free the GPU memory.
   * Used in CUDADeviceAPI::FreeWorkspace().
   *
   * \param ptr Pointer to the memory to be freed.
   */
  inline void CUDAFreeWorkspace(void* ptr) {
    auto entry = entrypoints_[Op::kCUDARawDelete];
    FUNCCAST(tensoradapter::CUDARawDelete, entry)(ptr);
  }

  /*!
  * \brief Find the current PyTorch CUDA stream
  * Used in runtime::getCurrentCUDAStream().
  * 
  * \note PyTorch pre-allocates/sets the current CUDA stream
  * on current device via cudaGetDevice(). Make sure to call cudaSetDevice()
  * before invoking this function.
  *
  * \return cudaStream_t stream handle
  */
  inline cudaStream_t CUDAGetCurrentStream() {
    auto entry = entrypoints_[Op::kCUDACurrentStream];
    return FUNCCAST(tensoradapter::CUDACurrentStream, entry)();
  }
#endif  // DGL_USE_CUDA

  /*!
   * \brief Record streams that are using this tensor.
   * Used in NDArray::RecordStream().
   *
   * \param ptr Pointer of the tensor to be recorded.
   * \param stream The stream that is using this tensor.
   * \param device_id Device of the tensor.
   */
  inline void RecordStream(void* ptr, DGLStreamHandle stream, int device_id) {
#ifdef DGL_USE_CUDA
    auto entry = entrypoints_[Op::kRecordStream];
    FUNCCAST(tensoradapter::RecordStream, entry)(
      ptr, static_cast<cudaStream_t>(stream), device_id);
#endif  // DGL_USE_CUDA
  }

 private:
  /*! \brief ctor */
  TensorDispatcher() = default;
  /*! \brief dtor */
  ~TensorDispatcher();

  /*!
   * \brief List of symbols in the adapter library.
   *
   * Must match the functions in tensoradapter/include/tensoradapter.h.
   */
  static constexpr const char *names_[] = {
    "CPURawAlloc",
    "CPURawDelete",
#ifdef DGL_USE_CUDA
    "CUDARawAlloc",
    "CUDARawDelete",
    "CUDACurrentStream",
    "RecordStream",
#endif  // DGL_USE_CUDA
  };

  /*! \brief Index of each function to the symbol list */
  class Op {
   public:
    static constexpr int kCPURawAlloc = 0;
    static constexpr int kCPURawDelete = 1;
#ifdef DGL_USE_CUDA
    static constexpr int kCUDARawAlloc = 2;
    static constexpr int kCUDARawDelete = 3;
    static constexpr int kCUDACurrentStream = 4;
    static constexpr int kRecordStream = 5;
#endif  // DGL_USE_CUDA
  };

  /*! \brief Number of functions */
  static constexpr int num_entries_ = sizeof(names_) / sizeof(names_[0]);

  /*! \brief Entrypoints of each function */
  void* entrypoints_[num_entries_] = {
    nullptr,
    nullptr,
#ifdef DGL_USE_CUDA
    nullptr,
    nullptr,
    nullptr,
    nullptr,
#endif  // DGL_USE_CUDA
  };

  bool available_ = false;
#if defined(WIN32) || defined(_WIN32)
  HINSTANCE handle_;
#else   // !WIN32
  void* handle_;
#endif  // WIN32
};

};  // namespace runtime
};  // namespace dgl

#undef FUNCCAST

#endif  // DGL_RUNTIME_TENSORDISPATCH_H_