MemoryOverlap.h 1.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. #pragma once
  2. #include <c10/macros/Export.h>
  3. namespace c10 {
  4. struct TensorImpl;
  5. }
  6. namespace at {
  7. class TensorBase;
  8. // MemOverlap: Whether or not there is memory overlap
  9. //
  10. // No: Absolutely no memory overlap
  11. // Yes: Absolutely yes memory overlap
  12. // TooHard: There might be memory overlap, but it was too expensive to compute.
  13. //
  14. // NB: Please update the python test for these if you renumber them.
  15. enum class MemOverlap { No, Yes, TooHard };
  16. enum class MemOverlapStatus { Full, Partial, No, TooHard };
  17. TORCH_API MemOverlap has_internal_overlap(const TensorBase& t);
  18. TORCH_API MemOverlap has_internal_overlap(c10::TensorImpl* t);
  19. TORCH_API void assert_no_internal_overlap(const TensorBase& t);
  20. TORCH_API void assert_no_internal_overlap(c10::TensorImpl* t);
  21. TORCH_API MemOverlapStatus
  22. get_overlap_status(const TensorBase& a, const TensorBase& b);
  23. TORCH_API MemOverlapStatus
  24. get_overlap_status(c10::TensorImpl* a, c10::TensorImpl* b);
  25. TORCH_API void assert_no_partial_overlap(
  26. const TensorBase& a,
  27. const TensorBase& b);
  28. void assert_no_partial_overlap(c10::TensorImpl* a, c10::TensorImpl* b);
  29. TORCH_API void assert_no_overlap(const TensorBase& a, const TensorBase& b);
  30. TORCH_API void assert_no_overlap(c10::TensorImpl* a, c10::TensorImpl* b);
  31. } // namespace at