jiterator.h 983 B

12345678910111213141516171819202122232425262728293031323334353637383940
  1. #pragma once
  2. #include <ATen/jit_macros.h>
  3. #if AT_USE_JITERATOR()
  4. #include <c10/macros/Export.h>
  5. #include <c10/util/SmallVector.h>
  6. #include <ATen/core/Tensor.h>
  7. #include <string>
  8. #include <vector>
  9. namespace at {
  10. namespace cuda {
  11. TORCH_CUDA_CPP_API c10::SmallVector<at::Tensor> CompileAndLaunchKernel(
  12. const std::string& code_string,
  13. const std::string& kernel_name,
  14. const int num_outputs,
  15. const c10::SmallVector<at::Tensor>& tensors,
  16. const c10::SmallVector<at::Scalar>& extra_args,
  17. bool return_by_ref);
  18. }} // namespace at::cuda
  19. #else
  20. namespace at { namespace cuda {
  21. TORCH_CUDA_CPP_API c10::SmallVector<at::Tensor> CompileAndLaunchKernel(
  22. const std::string& code_string,
  23. const std::string& kernel_name,
  24. const int num_outputs,
  25. const c10::SmallVector<at::Tensor>& tensors,
  26. const c10::SmallVector<at::Scalar>& extra_args,
  27. bool return_by_ref) {
  28. TORCH_CHECK(false, "Jiterator is not supported");
  29. }
  30. }} // namespace at::cuda
  31. #endif // AT_USE_JITERATOR()