Lower _conj_copy and alias operation. (#8686)

Author: ysiraichi

codegen/xla_native_functions.yaml

@@ -26,6 +26,7 @@ full_codegen:
   - clamp.Tensor
   - clamp_max.Tensor
   - clamp_min.Tensor
+  - _conj_copy
   - cos
   - cosh
   - elu
@@ -138,6 +139,7 @@ supported:
   - add.Scalar
   - add.Tensor
   - addmm
+  - alias
   - alias_copy
   - arange.start_out
   - as_strided_copy

test/cpp/test_aten_xla_tensor_2.cpp

@@ -2420,6 +2420,27 @@ TEST_F(AtenXlaTensorTest, TestAcoshInPlace) {
   ExpectCounterChanged("xla::acosh", cpp_test::GetIgnoredCounters());
 }
 
+TEST_F(AtenXlaTensorTest, TestAlias) {
+  torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
+  torch::Tensor b = torch::alias(a);
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = torch::alias(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/0e-5);
+  });
+}
+
+TEST_F(AtenXlaTensorTest, TestConj) {
+  torch::Tensor a =
+      torch::rand({2, 2}, torch::TensorOptions(torch::kComplexFloat));
+  torch::Tensor b = torch::conj(a);
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = torch::conj(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/0e-5);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestCos) {
   torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
   torch::Tensor b = torch::cos(a);

test/test_operations.py

@@ -2384,6 +2384,24 @@ def test_cummax_0_sized_dimension(self):
 
     self.assertEqual(actual, expected)
 
+  def test_conj(self):
+    # Leave the factory out of the fallback count.
+    tensor = torch.rand(2, 2, dtype=torch.complex64)
+
+    met.clear_all()
+
+    def run(device):
+      return torch.conj(tensor.to(device))
+
+    actual = run("cpu")
+    expected = run(xm.xla_device())
+
+    self.assertEqual(
+        met.executed_fallback_ops(), [],
+        message="expected no fallback operations.")
+    self.assertEqual(
+        actual, expected.cpu(), message="XLA results should match CPU results.")
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/aten_xla_type.cpp

@@ -863,12 +863,17 @@ at::Tensor XLANativeFunctions::addmm(const at::Tensor& self,
                             /*bias=*/bridge::GetXlaTensor(self)));
 }
 
-at::Tensor XLANativeFunctions::alias_copy(const at::Tensor& self) {
+at::Tensor XLANativeFunctions::alias(const at::Tensor& self) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
   return bridge::AtenFromXlaTensor(
       tensor_methods::alias(bridge::GetXlaTensor(self)));
 }
 
+at::Tensor XLANativeFunctions::alias_copy(const at::Tensor& self) {
+  TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
+  return alias(self);
+}
+
 at::Tensor& XLANativeFunctions::arange_out(const at::Scalar& start,
                                            const at::Scalar& end,
                                            const at::Scalar& step,
@@ -1333,8 +1338,12 @@ at::Tensor XLANativeFunctions::clone(
     const at::Tensor& self,
     std::optional<at::MemoryFormat> /* memory_format */) {
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
-  return bridge::AtenFromXlaTensor(
-      tensor_methods::clone(bridge::GetXlaTensor(self)));
+  auto tensor = bridge::GetXlaTensor(self);
+  if (self.is_conj()) {
+    // Materialize the conjugate if necessary.
+    tensor = tensor_methods::conj(tensor);
+  }
+  return bridge::AtenFromXlaTensor(tensor_methods::clone(tensor));
 }
 
 at::Tensor XLANativeFunctions::constant_pad_nd(const at::Tensor& self,

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -342,6 +342,11 @@ torch_xla::XlaOpVector ClampMinTensor::Lower(LoweringContext* loctx) const {
   return ReturnOp(xla::Max(xla_input, xla_other), loctx);
 }
 
+torch_xla::XlaOpVector ConjCopy::Lower(LoweringContext* loctx) const {
+  xla::XlaOp input = loctx->GetOutputOp(operand(0));
+  return ReturnOp(xla::Conj(input), loctx);
+}
+
 torch_xla::XlaOpVector Cos::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   if (xla::primitive_util::IsIntegralType(XlaHelpers::TypeOfXlaOp(xla_input))) {

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -424,6 +424,10 @@ xla::Shape ClampMinTensorOutputShape(const torch::lazy::Value& input,
                           lower_for_shape_fn);
 }
 
+xla::Shape ConjCopyOutputShape(const torch::lazy::Value& input) {
+  return GetXlaShape(input);
+}
+
 xla::Shape CosOutputShape(const torch::lazy::Value& input) {
   xla::Shape result_shape = GetXlaShape(input);
   if (xla::primitive_util::IsIntegralType(result_shape.element_type())) {

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -108,6 +108,8 @@ xla::Shape ClampMaxTensorOutputShape(const torch::lazy::Value& input,
 xla::Shape ClampMinTensorOutputShape(const torch::lazy::Value& input,
                                      const torch::lazy::Value& target);
 
+xla::Shape ConjCopyOutputShape(const torch::lazy::Value& input);
+
 xla::Shape CosOutputShape(const torch::lazy::Value& input);
 
 xla::Shape CoshOutputShape(const torch::lazy::Value& input);
@@ -287,4 +289,4 @@ xla::Shape TruncOutputShape(const torch::lazy::Value& input);
 
 }  // namespace torch_xla
 
-#endif  // XLA_TORCH_XLA_CSRC_OPS_OPS_XLA_SHAPE_FN_H_
+#endif  // XLA_TORCH_XLA_CSRC_OPS_OPS_XLA_SHAPE_FN_H_

torch_xla/csrc/tensor_methods.cpp

@@ -1238,6 +1238,11 @@ XLATensorPtr clone(const XLATensorPtr& input) {
   return cloned;
 }
 
+XLATensorPtr conj(const XLATensorPtr& input) {
+  auto ir = input->GetIrValue();
+  return input->CreateFrom(torch_xla::MakeNode<ConjCopy>(ir));
+}
+
 XLATensorPtr constant_pad_nd(const XLATensorPtr& input,
                              absl::Span<const int64_t> pad,
                              const at::Scalar& value) {

torch_xla/csrc/tensor_methods.h

@@ -327,6 +327,8 @@ XLATensorPtr clamp(const XLATensorPtr& input,
 
 XLATensorPtr clone(const XLATensorPtr& input);
 
+XLATensorPtr conj(const XLATensorPtr& input);
+
 // Pad with the given value and size specified by the given list of low and
 // high paddings.
 XLATensorPtr constant_pad_nd(const XLATensorPtr& input,