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