Add aten::ceil (#463)

- ceil.out
  - ceil
  - ceil_

---------

Co-authored-by: Feng Yuan <[email protected]>

Author: hjhee

src/ATen/native/xpu/UnaryOps.cpp

@@ -515,4 +515,40 @@ Tensor& XPUNativeFunctions::erfc_out(const Tensor& self, Tensor& out) {
   return out;
 }
 
+TensorIterator ceil_meta(const Tensor& self, Tensor& out) {
+  TORCH_CHECK(!self.is_complex(), "ceil is not supported for complex inputs");
+  TensorIterator iter;
+  iter.build_borrowing_unary_op(out, self);
+  return iter;
+}
+
+Tensor XPUNativeFunctions::ceil(const Tensor& self) {
+  if (c10::isIntegralType(self.scalar_type(), /*includeBool=*/false)) {
+    return self.clone();
+  }
+  Tensor out;
+  auto iter = ceil_meta(self, out);
+  native::xpu::ceil_kernel(iter);
+  return iter.output();
+}
+
+Tensor& XPUNativeFunctions::ceil_(Tensor& self) {
+  if (c10::isIntegralType(self.scalar_type(), /*includeBool=*/false)) {
+    return self;
+  }
+  auto iter = ceil_meta(self, self);
+  native::xpu::ceil_kernel(iter);
+  return self;
+}
+
+Tensor& XPUNativeFunctions::ceil_out(const Tensor& self, Tensor& out) {
+  if (c10::isIntegralType(self.scalar_type(), /*includeBool=*/false)) {
+    out.copy_(self);
+    return out;
+  }
+  auto iter = ceil_meta(self, out);
+  native::xpu::ceil_kernel(iter);
+  return out;
+}
+
 } // namespace at

src/ATen/native/xpu/XPUFallback.template

@@ -180,7 +180,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "bitwise_right_shift.Tensor_out",
     "cauchy_",
     "_cdist_backward",
-    "ceil.out",
     "channel_shuffle",
     "cholesky",
     "cholesky_inverse",

src/ATen/native/xpu/sycl/UnaryFractionKernels.cpp

@@ -55,4 +55,25 @@ void reciprocal_kernel(TensorIteratorBase& iter) {
       [&]() { gpu_kernel(iter, ReciprocalFunctor<scalar_t>()); });
 }
 
+template <typename scalar_t>
+struct CeilFunctor {
+  scalar_t operator()(const scalar_t a) const {
+    return std::ceil(a);
+  }
+};
+
+template <typename T>
+struct CeilFunctor<c10::complex<T>> {
+  c10::complex<T> operator()(const c10::complex<T> a) const {
+    return c10::complex<T>(std::ceil(a.real()), std::ceil(a.imag()));
+  }
+};
+
+void ceil_kernel(TensorIteratorBase& iter) {
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      ScalarType::Half, ScalarType::BFloat16, iter.dtype(), "ceil_xpu", [&]() {
+        gpu_kernel(iter, CeilFunctor<scalar_t>());
+      });
+}
+
 } // namespace at::native::xpu

src/ATen/native/xpu/sycl/UnaryFractionKernels.h

@@ -6,4 +6,6 @@ namespace at::native::xpu {
 
 void reciprocal_kernel(TensorIteratorBase& iter);
 
+void ceil_kernel(TensorIteratorBase& iter);
+
 } // namespace at::native::xpu

yaml/xpu_functions.yaml

@@ -500,3 +500,6 @@ supported:
   - randperm.generator_out
   - _amp_foreach_non_finite_check_and_unscale_
   - _amp_update_scale_
+  - ceil
+  - ceil_
+  - ceil.out