Add clamp impl back because XLANativeFunctions::hardtanh uses it

Author: ghpvnist

codegen/xla_native_functions.yaml

@@ -4,7 +4,7 @@
 # - https://github.com/pytorch/xla/blob/master/docs/source/contribute/codegen_migration.md
 backend: XLA
 cpp_namespace: torch_xla
-# full_codegen is the prefered method of code generation. Through this config
+# full_codegen is the preferred method of code generation. Through this config
 # ops get implementations (and IR classes) generated. See
 # https://github.com/pytorch/xla/blob/master/docs/source/contribute/codegen_migration.md
 # for more details on differences on what gets generated or not.

test/cpp/BUILD

@@ -130,21 +130,6 @@ ptxla_cc_test(
     ],
 )
 
-# ptxla_cc_test(
-#     name = "test_aten_xla_tensor_5",
-#     size = "enormous",
-#     srcs = ["test_aten_xla_tensor_5.cpp"],
-#     deps = [
-#         ":cpp_test_util",
-#         ":torch_xla_test",
-#         "//torch_xla/csrc/runtime:metrics",
-#         "//torch_xla/csrc:tensor",
-#         "//torch_xla/csrc:aten_cuda_functions",
-#         "@com_google_googletest//:gtest_main",
-#         "@xla//xla:permutation_util",
-#     ],
-# )
-
 # This tets is very large so it's split into shards.
 # To make it run fast, please add new shards when needed.
 [

torch_xla/csrc/ops/ops.cpp

@@ -195,6 +195,23 @@ torch::lazy::NodePtr SoftmaxBackwardOp(const torch::lazy::Value& grad_output,
           dim, GetXlaShape(grad_output).dimensions_size()));
 }
 
+torch::lazy::NodePtr Clamp(const torch::lazy::Value& input,
+                           const torch::lazy::Value& min,
+                           const torch::lazy::Value& max) {
+  auto lower_fn = [](const XlaNode& node,
+                     LoweringContext* loctx) -> XlaOpVector {
+    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
+    xla::XlaOp xla_min = loctx->GetOutputOp(node.operand(1));
+    xla::XlaOp xla_max = loctx->GetOutputOp(node.operand(2));
+    xla::PrimitiveType input_type = XlaHelpers::TypeOfXlaOp(xla_input);
+    xla_min = ConvertTo(xla_min, XlaHelpers::TypeOfXlaOp(xla_min), input_type);
+    xla_max = ConvertTo(xla_max, XlaHelpers::TypeOfXlaOp(xla_max), input_type);
+    return node.ReturnOp(xla::Clamp(xla_min, xla_input, xla_max), loctx);
+  };
+  return GenericOp(torch::lazy::OpKind(at::aten::clamp), {input, min, max},
+                   GetXlaShape(input), std::move(lower_fn));
+}
+
 torch::lazy::NodePtr Celu(const torch::lazy::Value& input,
                           const at::Scalar& alpha) {
   auto lower_fn = [=](const XlaNode& node,

torch_xla/csrc/ops/ops.h

@@ -121,6 +121,10 @@ torch::lazy::NodePtr SoftmaxBackwardOp(const torch::lazy::Value& grad_output,
                                        const torch::lazy::Value& output,
                                        int64_t dim);
 
+torch::lazy::NodePtr Clamp(const torch::lazy::Value& input,
+                           const torch::lazy::Value& min,
+                           const torch::lazy::Value& max);
+
 torch::lazy::NodePtr Celu(const torch::lazy::Value& input,
                           const at::Scalar& alpha);
 

torch_xla/csrc/tensor_methods.cpp

@@ -1226,6 +1226,14 @@ void celu_(XLATensorPtr& input, const at::Scalar& alpha) {
   input->SetInPlaceIrValue(Celu(input->GetIrValue(), alpha));
 }
 
+XLATensorPtr clamp(const XLATensorPtr& input,
+                   const std::optional<at::Scalar>& min,
+                   const std::optional<at::Scalar>& max) {
+  MinMaxValues min_max = GetMinMaxValues(input, min, max);
+  return input->CreateFrom(
+      Clamp(input->GetIrValue(), min_max.min, min_max.max));
+}
+
 XLATensorPtr clone(const XLATensorPtr& input) {
   XLATensorPtr cloned = input->CreateFrom(input->GetIrValue());
   if (input->sharding_spec() != nullptr) {

torch_xla/csrc/tensor_methods.h

@@ -320,6 +320,13 @@ XLATensorPtr pixel_shuffle(const XLATensorPtr& self, int64_t upscale_factor);
 XLATensorPtr celu(const XLATensorPtr& input, const at::Scalar& alpha);
 void celu_(XLATensorPtr& input, const at::Scalar& alpha);
 
+XLATensorPtr clamp(const XLATensorPtr& input,
+                   const std::optional<at::Scalar>& min,
+                   const std::optional<at::Scalar>& max);
+XLATensorPtr clamp(const XLATensorPtr& input,
+                   const std::optional<at::Tensor>& min,
+                   const std::optional<at::Tensor>& max);
+
 XLATensorPtr clone(const XLATensorPtr& input);
 
 XLATensorPtr conj(const XLATensorPtr& input);