Add aten::_foreach_clamp_max (#967)

- `_foreach_clamp_max.List` - `_foreach_clamp_max_.List` - `_foreach_clamp_max.Scalar` - `_foreach_clamp_max_.Scalar` - `_foreach_clamp_max.ScalarList` - `_foreach_clamp_max_.ScalarList` --------- Co-authored-by: min.jean.cho <[email protected]> Co-authored-by: Yutao Xu <[email protected]>
intel · Oct 18, 2024 · 9fc5d81 · 9fc5d81
1 parent f592ff9
commit 9fc5d81
Show file tree

Hide file tree

Showing 10 changed files with 81 additions and 0 deletions.
diff --git a/src/ATen/native/xpu/ForeachOpList.cpp b/src/ATen/native/xpu/ForeachOpList.cpp
@@ -2,6 +2,7 @@
 #include <ATen/ops/_foreach_add_native.h>
 #include <ATen/ops/_foreach_addcdiv_native.h>
 #include <ATen/ops/_foreach_addcmul_native.h>
+#include <ATen/ops/_foreach_clamp_max_native.h>
 #include <ATen/ops/_foreach_div_native.h>
 #include <ATen/ops/_foreach_lerp_native.h>
 #include <ATen/ops/_foreach_mul_native.h>
@@ -65,6 +66,7 @@ namespace native {
 FOREACH_BINARY_OP_LIST_ALPHA(add);
 FOREACH_BINARY_OP_LIST(mul, false);
 FOREACH_BINARY_OP_LIST(div, true);
+FOREACH_BINARY_OP_LIST(clamp_max, true);
 FOREACH_BINARY_OP_LIST(clamp_min, true);
 
 #define FOREACH_POINTWISE_OP_TENSOR(NAME)                                  \

diff --git a/src/ATen/native/xpu/ForeachOpScalar.cpp b/src/ATen/native/xpu/ForeachOpScalar.cpp
@@ -2,6 +2,7 @@
 #include <ATen/ops/_foreach_add_native.h>
 #include <ATen/ops/_foreach_addcdiv_native.h>
 #include <ATen/ops/_foreach_addcmul_native.h>
+#include <ATen/ops/_foreach_clamp_max_native.h>
 #include <ATen/ops/_foreach_div_native.h>
 #include <ATen/ops/_foreach_lerp_native.h>
 #include <ATen/ops/_foreach_mul_native.h>
@@ -38,6 +39,7 @@ namespace native {
 FOREACH_BINARY_OP_SCALAR(add, /*div_op*/ false);
 FOREACH_BINARY_OP_SCALAR(mul, /*div_op*/ false);
 FOREACH_BINARY_OP_SCALAR(div, /*div_op*/ true);
+FOREACH_BINARY_OP_SCALAR(clamp_max, /*div_op*/ true);
 FOREACH_BINARY_OP_SCALAR(clamp_min, /*div_op*/ true);
 
 #define FOREACH_POINTWISE_OP_SCALAR(NAME)                                   \

diff --git a/src/ATen/native/xpu/ForeachOpScalarList.cpp b/src/ATen/native/xpu/ForeachOpScalarList.cpp
@@ -2,6 +2,7 @@
 #include <ATen/ops/_foreach_add_native.h>
 #include <ATen/ops/_foreach_addcdiv_native.h>
 #include <ATen/ops/_foreach_addcmul_native.h>
+#include <ATen/ops/_foreach_clamp_max_native.h>
 #include <ATen/ops/_foreach_div_native.h>
 #include <ATen/ops/_foreach_mul_native.h>
 #include <ATen/ops/_foreach_clamp_min_native.h>
@@ -41,6 +42,7 @@ namespace native {
 FOREACH_BINARY_OP_SCALARLIST(add, /*div_op*/ false);
 FOREACH_BINARY_OP_SCALARLIST(mul, /*div_op*/ false);
 FOREACH_BINARY_OP_SCALARLIST(div, /*div_op*/ true);
+FOREACH_BINARY_OP_SCALARLIST(clamp_max, /*div_op*/ true);
 FOREACH_BINARY_OP_SCALARLIST(clamp_min, /*div_op*/ true);
 
 #define FOREACH_POINTWISE_OP_SCALARLIST(NAME)                                \

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpListKernels.cpp b/src/ATen/native/xpu/sycl/ForeachBinaryOpListKernels.cpp
@@ -182,6 +182,14 @@ FOREACH_BINARY_LIST_KERNEL(div) {
   return all_types_complex_bool_half_bfloat16<std::divides>(tensor1, tensor2);
 }
 
+FOREACH_BINARY_LIST_INPLACE_KERNEL(clamp_max) {
+  return all_types_half_bfloat16_<foreach_internal::minimum>(tensor1, tensor2);
+}
+
+FOREACH_BINARY_LIST_KERNEL(clamp_max) {
+  return all_types_half_bfloat16<foreach_internal::minimum>(tensor1, tensor2);
+}
+
 FOREACH_BINARY_LIST_INPLACE_KERNEL(clamp_min) {
   return all_types_half_bfloat16_<foreach_internal::maximum>(tensor1, tensor2);
 }

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpListKernels.h b/src/ATen/native/xpu/sycl/ForeachBinaryOpListKernels.h
@@ -37,6 +37,8 @@ TORCH_XPU_API FOREACH_BINARY_LIST_INPLACE_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_LIST_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_LIST_INPLACE_KERNEL(div);
 TORCH_XPU_API FOREACH_BINARY_LIST_KERNEL(div);
+TORCH_XPU_API FOREACH_BINARY_LIST_INPLACE_KERNEL(clamp_max);
+TORCH_XPU_API FOREACH_BINARY_LIST_KERNEL(clamp_max);
 TORCH_XPU_API FOREACH_BINARY_LIST_INPLACE_KERNEL(clamp_min);
 TORCH_XPU_API FOREACH_BINARY_LIST_KERNEL(clamp_min);
 

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarKernels.cpp b/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarKernels.cpp
@@ -150,6 +150,14 @@ FOREACH_BINARY_SCALAR_KERNEL(div) {
   return all_types_complex_bool_half_bfloat16<std::divides>(tensors, scalar);
 }
 
+FOREACH_BINARY_SCALAR_INPLACE_KERNEL(clamp_max) {
+  return all_types_half_bfloat16_<foreach_internal::minimum>(tensors, scalar);
+}
+
+FOREACH_BINARY_SCALAR_KERNEL(clamp_max) {
+  return all_types_half_bfloat16<foreach_internal::minimum>(tensors, scalar);
+}
+
 FOREACH_BINARY_SCALAR_INPLACE_KERNEL(clamp_min) {
   return all_types_half_bfloat16_<foreach_internal::maximum>(tensors, scalar);
 }

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarKernels.h b/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarKernels.h
@@ -23,6 +23,8 @@ TORCH_XPU_API FOREACH_BINARY_SCALAR_INPLACE_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_SCALAR_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_SCALAR_INPLACE_KERNEL(div);
 TORCH_XPU_API FOREACH_BINARY_SCALAR_KERNEL(div);
+TORCH_XPU_API FOREACH_BINARY_SCALAR_INPLACE_KERNEL(clamp_max);
+TORCH_XPU_API FOREACH_BINARY_SCALAR_KERNEL(clamp_max);
 TORCH_XPU_API FOREACH_BINARY_SCALAR_INPLACE_KERNEL(clamp_min);
 TORCH_XPU_API FOREACH_BINARY_SCALAR_KERNEL(clamp_min);
 

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarListKernels.cpp b/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarListKernels.cpp
@@ -155,6 +155,14 @@ FOREACH_BINARY_SCALARLIST_KERNEL(div) {
   return all_types_complex_bool_half_bfloat16<std::divides>(tensors, scalars);
 }
 
+FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(clamp_max) {
+  return all_types_half_bfloat16_<foreach_internal::minimum>(tensors, scalars);
+}
+
+FOREACH_BINARY_SCALARLIST_KERNEL(clamp_max) {
+  return all_types_half_bfloat16<foreach_internal::minimum>(tensors, scalars);
+}
+
 FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(clamp_min) {
   return all_types_half_bfloat16_<foreach_internal::maximum>(tensors, scalars);
 }

diff --git a/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarListKernels.h b/src/ATen/native/xpu/sycl/ForeachBinaryOpScalarListKernels.h
@@ -23,6 +23,8 @@ TORCH_XPU_API FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_SCALARLIST_KERNEL(mul);
 TORCH_XPU_API FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(div);
 TORCH_XPU_API FOREACH_BINARY_SCALARLIST_KERNEL(div);
+TORCH_XPU_API FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(clamp_max);
+TORCH_XPU_API FOREACH_BINARY_SCALARLIST_KERNEL(clamp_max);
 TORCH_XPU_API FOREACH_BINARY_SCALARLIST_INPLACE_KERNEL(clamp_min);
 TORCH_XPU_API FOREACH_BINARY_SCALARLIST_KERNEL(clamp_min);
 

diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml
@@ -2003,6 +2003,51 @@
     XPU: foreach_tensor_div_scalar_kernel_xpu_
   autogen: _foreach_div.Scalar_out
 
+- func: _foreach_clamp_max.Scalar(Tensor[] self, Scalar scalar) -> Tensor[]
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow
+    XPU: foreach_tensor_clamp_max_scalar_kernel_xpu
+
+- func: _foreach_clamp_max_.Scalar(Tensor(a!)[] self, Scalar scalar) -> ()
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow_
+    XPU: foreach_tensor_clamp_max_scalar_kernel_xpu_
+  autogen: _foreach_clamp_max.Scalar_out
+
+- func: _foreach_clamp_max.List(Tensor[] self, Tensor[] other) -> Tensor[]
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow
+    XPU: foreach_tensor_clamp_max_list_kernel_xpu
+
+- func: _foreach_clamp_max_.List(Tensor(a!)[] self, Tensor[] other) -> ()
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow_
+    XPU: foreach_tensor_clamp_max_list_kernel_xpu_
+  autogen: _foreach_clamp_max.List_out
+
+- func: _foreach_clamp_max.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[]
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow
+    XPU: foreach_tensor_clamp_max_scalarlist_kernel_xpu
+
+- func: _foreach_clamp_max_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> ()
+  device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
+  variants: function
+  dispatch:
+    CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow_
+    XPU: foreach_tensor_clamp_max_scalarlist_kernel_xpu_
+  autogen: _foreach_clamp_max.ScalarList_out
+
 - func: _foreach_addcmul.Scalar(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar value=1) -> Tensor[]
   device_check: NoCheck   # foreach kernels fall back to slow path when tensor are on different devices
   variants: function