CPUFixedAllocator.h 845 B

123456789101112131415161718192021222324252627282930313233
  1. #pragma once
  2. #include <c10/core/Allocator.h>
  3. #include <c10/util/Exception.h>
  4. // This file creates a fake allocator that just throws exceptions if
  5. // it is actually used.
  6. // state passed to the allocator is the std::function<void(void*)> called
  7. // when the blob is release by ATen
  8. namespace at {
  9. static cpu_fixed_malloc(void*, ptrdiff_t) {
  10. AT_ERROR("attempting to resize a tensor view of an external blob");
  11. }
  12. static cpu_fixed_realloc(void*, void*, ptrdiff_t) {
  13. AT_ERROR("attempting to resize a tensor view of an external blob");
  14. }
  15. static cpu_fixed_free(void* state, void* allocation) {
  16. auto on_release = static_cast<std::function<void(void*)>*>(state);
  17. (*on_release)(allocation);
  18. delete on_release;
  19. }
  20. static Allocator CPU_fixed_allocator = {
  21. cpu_fixed_malloc,
  22. cpu_fixed_realloc,
  23. cpu_fixed_free};
  24. } // namespace at