GridSamplerKernel.h 838 B

12345678910111213141516171819202122232425262728293031323334
  1. #pragma once
  2. #include <ATen/native/DispatchStub.h>
  3. #include <array>
  4. #include <cstdint>
  5. namespace at {
  6. class TensorBase;
  7. }
  8. namespace at { namespace native {
  9. using forward_2d_fn = void (*) (
  10. const TensorBase &output,
  11. const TensorBase &input,
  12. const TensorBase &grid,
  13. int64_t interpolation_mode,
  14. int64_t padding_mode,
  15. bool align_corners);
  16. using backward_2d_fn = void (*) (
  17. const TensorBase &grad_input,
  18. const TensorBase &grad_grid,
  19. const TensorBase &grad_output,
  20. const TensorBase &input,
  21. const TensorBase &grid,
  22. int64_t interpolation_mode,
  23. int64_t padding_mode,
  24. bool align_corners,
  25. std::array<bool, 2> output_mask);
  26. DECLARE_DISPATCH(forward_2d_fn, grid_sampler_2d_cpu_kernel);
  27. DECLARE_DISPATCH(backward_2d_fn, grid_sampler_2d_backward_cpu_kernel);
  28. }} // namespace at::native