ScanKernels.h 779 B

123456789101112131415161718
  1. #pragma once
  2. #include <cstdint>
  3. namespace at {
  4. class TensorBase;
  5. namespace native {
  6. // NOTE: these functions require output tensors to be contiguous
  7. void launch_cummax_cuda_kernel(const TensorBase& self, const TensorBase& values,
  8. const TensorBase& indices, int64_t dim);
  9. void launch_cummin_cuda_kernel(const TensorBase& self, const TensorBase& values,
  10. const TensorBase& indices, int64_t dim);
  11. void launch_logcumsumexp_cuda_kernel(const TensorBase& result, const TensorBase& self, int64_t dim);
  12. void launch_cumsum_cuda_kernel(const TensorBase& result, const TensorBase& self, int64_t dim);
  13. void launch_cumprod_cuda_kernel(const TensorBase& result, const TensorBase& self, int64_t dim);
  14. }} // namespace at::native