Program Listing for File NPUHooks.h

↰ Return to documentation for file (csrc/backend/NPUHooks.h)

#pragma once
#include <ATen/detail/PrivateUse1HooksInterface.h>
#include <c10/core/Storage.h>
#include "csrc/backend/NPUGeneratorImpl.h"

namespace c10::backend {

struct TORCH_API NPUHooks : public at::PrivateUse1HooksInterface {
  virtual ~NPUHooks() = default;
  const at::Generator& getDefaultGenerator(
      c10::DeviceIndex device_index) override {
    static auto device_gen =
        c10::backend::detail::getDefaultNPUGenerator(device_index);
    return device_gen;
  }
  void initPrivateUse1() const override;
  bool hasPrimaryContext(c10::DeviceIndex device_index) const override;
  void resizePrivateUse1Bytes(const c10::Storage& storage, size_t new_bytes)
      const override;
  bool isPinnedPtr(const void* data) const;
  at::Allocator* getPinnedMemoryAllocator() const override;
};

struct TORCH_API NPUHooksArgs : public at::PrivateUse1HooksArgs {};

// register to PrivateUse1HooksInterface
at::PrivateUse1HooksInterface* get_npu_hooks();
} // namespace c10::npu