WIP

Author: ghpvnist

codegen/xla_native_functions.yaml

@@ -35,7 +35,9 @@ full_codegen:
   - cholesky
   - clamp
   - clamp.Tensor
+  - clamp_max
   - clamp_max.Tensor
+  - clamp_min
   - clamp_min.Tensor
   - _conj_copy
   - cos
@@ -184,8 +186,6 @@ supported:
   - cat
   - celu
   - celu_
-  - clamp_max
-  - clamp_min
   - clone
   - constant_pad_nd
   - convolution_backward_overrideable

test/cpp/BUILD

@@ -130,22 +130,37 @@ 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.
-[
-    ptxla_cc_test(
-        name = test[:-len(".cpp")],
-        size = "enormous",
-        srcs = [test],
-        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",
-        ],
-    )
-    for test in glob(["test_aten_xla_tensor*cpp"])
-]
+# [
+#     ptxla_cc_test(
+#         name = test[:-len(".cpp")],
+#         size = "enormous",
+#         srcs = [test],
+#         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",
+#         ],
+#     )
+#     for test in glob(["test_aten_xla_tensor*cpp"])
+# ]

torch_xla/csrc/aten_xla_type.cpp

@@ -1312,20 +1312,6 @@ at::Tensor& XLANativeFunctions::celu_(at::Tensor& self,
   return self;
 }
 
-at::Tensor XLANativeFunctions::clamp_max(const at::Tensor& self,
-                                         const at::Scalar& max) {
-  TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  return bridge::AtenFromXlaTensor(
-      tensor_methods::clamp(bridge::GetXlaTensor(self), std::nullopt, max));
-}
-
-at::Tensor XLANativeFunctions::clamp_min(const at::Tensor& self,
-                                         const at::Scalar& min) {
-  TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  return bridge::AtenFromXlaTensor(
-      tensor_methods::clamp(bridge::GetXlaTensor(self), min, std::nullopt));
-}
-
 at::Tensor XLANativeFunctions::clone(
     const at::Tensor& self,
     std::optional<at::MemoryFormat> /* memory_format */) {

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -401,12 +401,24 @@ torch_xla::XlaOpVector ClampTensor::Lower(LoweringContext* loctx) const {
   return ReturnOp(res, loctx);
 }
 
+torch_xla::XlaOpVector ClampMax::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_other = loctx->GetOutputOp(operand(1));
+  return ReturnOp(xla::Min(xla_input, xla_other), loctx);
+}
+
 torch_xla::XlaOpVector ClampMaxTensor::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   xla::XlaOp xla_other = loctx->GetOutputOp(operand(1));
   return ReturnOp(xla::Min(xla_input, xla_other), loctx);
 }
 
+torch_xla::XlaOpVector ClampMin::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_other = loctx->GetOutputOp(operand(1));
+  return ReturnOp(xla::Max(xla_input, xla_other), loctx);
+}
+
 torch_xla::XlaOpVector ClampMinTensor::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   xla::XlaOp xla_other = loctx->GetOutputOp(operand(1));

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -454,6 +454,16 @@ xla::Shape ClampTensorOutputShape(
   return InferOutputShape(shapes, lower_for_shape_fn);
 }
 
+xla::Shape ClampMaxOutputShape(const torch::lazy::Value& input,
+                               const torch::lazy::Value& target) {
+  auto lower_for_shape_fn =
+      [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
+    return xla::Min(operands[0], operands[1]);
+  };
+  return InferOutputShape({GetXlaShape(input), GetXlaShape(other)},
+                          lower_for_shape_fn);
+}
+
 xla::Shape ClampMaxTensorOutputShape(const torch::lazy::Value& input,
                                      const torch::lazy::Value& other) {
   auto lower_for_shape_fn =
@@ -464,6 +474,16 @@ xla::Shape ClampMaxTensorOutputShape(const torch::lazy::Value& input,
                           lower_for_shape_fn);
 }
 
+xla::Shape ClampMinOutputShape(const torch::lazy::Value& input,
+                               const torch::lazy::Value& other) {
+  auto lower_for_shape_fn =
+      [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
+    return xla::Max(operands[0], operands[1]);
+  };
+  return InferOutputShape({GetXlaShape(input), GetXlaShape(other)},
+                          lower_for_shape_fn);
+}
+
 xla::Shape ClampMinTensorOutputShape(const torch::lazy::Value& input,
                                      const torch::lazy::Value& other) {
   auto lower_for_shape_fn =

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -112,9 +112,15 @@ xla::Shape ClampTensorOutputShape(const torch::lazy::Value& input,
                                   const std::optional<torch::lazy::Value>& min,
                                   const std::optional<torch::lazy::Value>& max);
 
+xla::Shape ClampMaxOutputShape(const torch::lazy::Value& input,
+                               const torch::lazy::Value& target);
+
 xla::Shape ClampMaxTensorOutputShape(const torch::lazy::Value& input,
                                      const torch::lazy::Value& target);
 
+xla::Shape ClampMinOutputShape(const torch::lazy::Value& input,
+                               const torch::lazy::Value& target);
+
 xla::Shape ClampMinTensorOutputShape(const torch::lazy::Value& input,
                                      const torch::lazy::Value& target);