CUDAGraphsUtils.cuh 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. #pragma once
  2. #include <ATen/cuda/CUDAGeneratorImpl.h>
  3. #include <ATen/cuda/CUDAEvent.h>
  4. #include <ATen/cuda/PhiloxUtils.cuh>
  5. #include <ATen/cuda/detail/CUDAHooks.h>
  6. #include <ATen/detail/CUDAHooksInterface.h>
  7. #include <c10/core/StreamGuard.h>
  8. #include <c10/cuda/CUDAGraphsC10Utils.h>
  9. #include <c10/cuda/CUDAGuard.h>
  10. // c10/cuda/CUDAGraphsC10Utils.h has utils used by both c10 and aten.
  11. // This file adds utils used by aten only.
  12. namespace at::cuda {
  13. using CaptureId_t = c10::cuda::CaptureId_t;
  14. using CaptureStatus = c10::cuda::CaptureStatus;
  15. // Use this version where you don't want to create a CUDA context if none exists.
  16. inline CaptureStatus currentStreamCaptureStatus() {
  17. // don't create a context if we don't have to
  18. if (c10::cuda::hasPrimaryContext(c10::cuda::current_device())) {
  19. return c10::cuda::currentStreamCaptureStatusMayInitCtx();
  20. } else {
  21. return CaptureStatus::None;
  22. }
  23. }
  24. inline void assertNotCapturing(const std::string& attempt) {
  25. auto status = currentStreamCaptureStatus();
  26. TORCH_CHECK(status == CaptureStatus::None,
  27. attempt,
  28. " during CUDA graph capture. If you need this call to be captured, "
  29. "please file an issue. "
  30. "Current cudaStreamCaptureStatus: ",
  31. status);
  32. }
  33. inline void errorIfCapturingCudnnBenchmark(const std::string& version_specific) {
  34. auto status = currentStreamCaptureStatus();
  35. TORCH_CHECK(status == CaptureStatus::None,
  36. "Current cudaStreamCaptureStatus: ",
  37. status,
  38. "\nCapturing ",
  39. version_specific,
  40. "is prohibited. Possible causes of this error:\n"
  41. "1. No warmup iterations occurred before capture.\n"
  42. "2. The convolutions you're trying to capture use dynamic shapes, "
  43. "in which case capturing them is generally prohibited.");
  44. }
  45. } // namespace at::cuda