MemoryOverlap.h 1.5 KB

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