Migrate cosh to structured kernel (#55563)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/55563
Test Plan: Imported from OSS
Reviewed By: pbelevich
Differential Revision: D27653767
Pulled By: SplitInfinity
fbshipit-source-id: 11cd631679b9b5a88443a714a56f4178f5bf41b0
diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp
index ac270c1..402652b 100644
--- a/aten/src/ATen/native/UnaryOps.cpp
+++ b/aten/src/ATen/native/UnaryOps.cpp
@@ -37,6 +37,7 @@
CREATE_UNARY_META_FUNC(sin)
CREATE_UNARY_META_FUNC(sinc)
CREATE_UNARY_META_FUNC(sinh)
+CREATE_UNARY_META_FUNC(cosh)
CREATE_UNARY_META_FUNC(acosh)
} // namespace meta
@@ -415,11 +416,7 @@
CREATE_UNARY_TORCH_IMPL_FUNC(sinc)
CREATE_UNARY_TORCH_IMPL_FUNC(sinh)
-
-Tensor& cosh_out(const Tensor& self, Tensor& result) { return unary_op_impl_float_out(result, self, cosh_stub); }
-Tensor cosh(const Tensor& self) { return unary_op_impl_float(self, cosh_stub); }
-Tensor& cosh_(Tensor& self) { return unary_op_impl_(self, at::cosh_out); }
-
+CREATE_UNARY_TORCH_IMPL_FUNC(cosh)
CREATE_UNARY_TORCH_IMPL_FUNC(acosh)
// arccosh, alias for acosh
diff --git a/aten/src/ATen/native/UnaryOps.h b/aten/src/ATen/native/UnaryOps.h
index 2f50b3e..f870d2d 100644
--- a/aten/src/ATen/native/UnaryOps.h
+++ b/aten/src/ATen/native/UnaryOps.h
@@ -30,7 +30,7 @@
DECLARE_DISPATCH(unary_fn_with_scalar, clamp_max_stub);
DECLARE_DISPATCH(unary_fn_with_scalar, clamp_min_stub);
DECLARE_DISPATCH(unary_fn, cos_stub);
-DECLARE_DISPATCH(unary_fn, cosh_stub);
+DECLARE_DISPATCH(structured_unary_fn, cosh_stub);
DECLARE_DISPATCH(unary_fn, digamma_stub);
DECLARE_DISPATCH(unary_fn, entr_stub);
DECLARE_DISPATCH(unary_fn, erf_stub);
diff --git a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
index 463ffa6..b66b545 100644
--- a/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
+++ b/aten/src/ATen/native/cpu/UnaryOpsKernel.cpp
@@ -327,7 +327,7 @@
});
}
-static void cosh_kernel(TensorIterator& iter) {
+static void cosh_kernel(TensorIteratorBase& iter) {
AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(iter.dtype(), "cosh_cpu", [&]() {
cpu_kernel_vec(
iter,
diff --git a/aten/src/ATen/native/cuda/UnaryGeometricKernels.cu b/aten/src/ATen/native/cuda/UnaryGeometricKernels.cu
index 28fb5e4..4ac658e 100644
--- a/aten/src/ATen/native/cuda/UnaryGeometricKernels.cu
+++ b/aten/src/ATen/native/cuda/UnaryGeometricKernels.cu
@@ -58,7 +58,7 @@
});
}
-void cosh_kernel_cuda(TensorIterator& iter) {
+void cosh_kernel_cuda(TensorIteratorBase& iter) {
AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND1(ScalarType::Half, iter.common_dtype(), "cosh_cuda", [&]() {
gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
return ::cosh(a);
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index 4a96aed..5441719 100644
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -1063,15 +1063,15 @@
- func: cosh(Tensor self) -> Tensor
variants: function, method
- dispatch:
- CPU, CUDA: cosh
+ structured_delegate: cosh.out
- func: cosh_(Tensor(a!) self) -> Tensor(a!)
variants: function, method
- dispatch:
- CompositeExplicitAutograd: cosh_
+ structured_delegate: cosh.out
- func: cosh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
+ structured: True
+ structured_inherits: TensorIteratorBase
dispatch:
CPU, CUDA: cosh_out
