Add trunc to ATen XLA tensor

Author: asuhan

test/cpp/test_aten_xla_tensor.cpp

@@ -1143,6 +1143,16 @@ TEST_F(AtenXlaTensorTest, TestFloor) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestTrunc) {
+  at::Tensor a = at::randn({2, 2}, at::TensorOptions(at::kFloat)) * 100.0;
+  at::Tensor b = at::trunc(a);
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_a = bridge::CreateXlaTensor(a, device);
+    at::Tensor xla_b = at::trunc(xla_a);
+    AllClose(b, xla_b);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestNeg) {
   at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
   at::Tensor b = at::neg(a);

torch_xla/csrc/aten_xla_type.cpp

@@ -1027,6 +1027,17 @@ at::Tensor& AtenXlaType::floor_(at::Tensor& self) const {
   return self;
 }
 
+at::Tensor AtenXlaType::trunc(const at::Tensor& self) const {
+  return bridge::AtenFromXlaTensor(
+      XLATensor::trunc(bridge::GetXlaTensor(self)));
+}
+
+at::Tensor& AtenXlaType::trunc_(at::Tensor& self) const {
+  XLATensor self_tensor = bridge::GetXlaTensor(self);
+  XLATensor::trunc_(self_tensor);
+  return self;
+}
+
 int64_t AtenXlaType::size(const at::Tensor& self, int64_t dim) const {
   return bridge::GetXlaTensor(self).size(dim);
 }

torch_xla/csrc/aten_xla_type.h

@@ -300,6 +300,9 @@ class AtenXlaType : public AtenXlaTypeBase {
   at::Tensor floor(const at::Tensor& self) const override;
   at::Tensor& floor_(at::Tensor& self) const override;
 
+  at::Tensor trunc(const at::Tensor& self) const override;
+  at::Tensor& trunc_(at::Tensor& self) const override;
+
   int64_t size(const at::Tensor& self, int64_t dim) const override;
 
   std::tuple<at::Tensor, at::Tensor> kthvalue(const at::Tensor& self, int64_t k,

torch_xla/csrc/ops/arithmetic_ir_ops.cpp

@@ -9,7 +9,7 @@
 namespace torch_xla {
 namespace ir {
 
-Value operator+(const Value& node1, const Value& node2) {
+NodePtr operator+(const Value& node1, const Value& node2) {
   auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
     xla::XlaOp op0 = loctx->GetOutputOp(node.operand(0));
     xla::XlaOp op1 = loctx->GetOutputOp(node.operand(1));
@@ -21,7 +21,7 @@ Value operator+(const Value& node1, const Value& node2) {
       std::move(lower_fn));
 }
 
-Value operator-(const Value& node1, const Value& node2) {
+NodePtr operator-(const Value& node1, const Value& node2) {
   auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
     xla::XlaOp op0 = loctx->GetOutputOp(node.operand(0));
     xla::XlaOp op1 = loctx->GetOutputOp(node.operand(1));
@@ -33,7 +33,7 @@ Value operator-(const Value& node1, const Value& node2) {
       std::move(lower_fn));
 }
 
-Value operator*(const Value& node1, const Value& node2) {
+NodePtr operator*(const Value& node1, const Value& node2) {
   auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
     xla::XlaOp op0 = loctx->GetOutputOp(node.operand(0));
     xla::XlaOp op1 = loctx->GetOutputOp(node.operand(1));
@@ -45,7 +45,7 @@ Value operator*(const Value& node1, const Value& node2) {
       std::move(lower_fn));
 }
 
-Value operator/(const Value& node1, const Value& node2) {
+NodePtr operator/(const Value& node1, const Value& node2) {
   auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
     xla::XlaOp op0 = loctx->GetOutputOp(node.operand(0));
     xla::XlaOp op1 = loctx->GetOutputOp(node.operand(1));

torch_xla/csrc/ops/arithmetic_ir_ops.h

@@ -5,10 +5,10 @@
 namespace torch_xla {
 namespace ir {
 
-Value operator+(const Value& node1, const Value& node2);
-Value operator-(const Value& node1, const Value& node2);
-Value operator*(const Value& node1, const Value& node2);
-Value operator/(const Value& node1, const Value& node2);
+NodePtr operator+(const Value& node1, const Value& node2);
+NodePtr operator-(const Value& node1, const Value& node2);
+NodePtr operator*(const Value& node1, const Value& node2);
+NodePtr operator/(const Value& node1, const Value& node2);
 
 }  // namespace ir
 }  // namespace torch_xla

torch_xla/csrc/ops/ops.cpp

@@ -11,6 +11,7 @@
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/lowering_context.h"
 #include "torch_xla/csrc/nll_loss.h"
+#include "torch_xla/csrc/ops/arithmetic_ir_ops.h"
 #include "torch_xla/csrc/ops/constant.h"
 #include "torch_xla/csrc/ops/infer_output_shape.h"
 #include "torch_xla/csrc/pooling.h"
@@ -73,6 +74,8 @@ PTXLA_BINARY_OP(Pow, at::aten::pow, xla::Pow);
 PTXLA_BINARY_OP(Fmod, at::aten::fmod, xla::Rem);
 PTXLA_BINARY_OP(Atan2, at::aten::atan2, xla::Atan2);
 
+NodePtr Trunc(const Value& input) { return Floor(Abs(input)) * SignOp(input); }
+
 NodePtr LogBase(const Value& input, OpKind op, double base) {
   auto lower_fn = [base](const Node& node,
                          LoweringContext* loctx) -> XlaOpVector {

torch_xla/csrc/ops/ops.h

@@ -114,6 +114,8 @@ NodePtr Ceil(const Value& input);
 
 NodePtr Floor(const Value& input);
 
+NodePtr Trunc(const Value& input);
+
 NodePtr AddMatMulOp(const Value& input, const Value& weight, const Value& bias);
 
 NodePtr Dot(const Value& input, const Value& weight);

torch_xla/csrc/tensor.cpp

@@ -1177,6 +1177,14 @@ void XLATensor::floor_(XLATensor& input) {
   input.SetIrValue(ir::ops::Floor(input.GetIrValue()));
 }
 
+XLATensor XLATensor::trunc(const XLATensor& input) {
+  return input.CreateFrom(ir::ops::Trunc(input.GetIrValue()));
+}
+
+void XLATensor::trunc_(XLATensor& input) {
+  input.SetIrValue(ir::ops::Trunc(input.GetIrValue()));
+}
+
 XLATensor XLATensor::slice(const XLATensor& input, xla::int64 dim,
                            xla::int64 start, xla::int64 end, xla::int64 step) {
   auto input_shape = input.shape();

torch_xla/csrc/tensor.h

@@ -359,6 +359,9 @@ class XLATensor {
   static XLATensor floor(const XLATensor& input);
   static void floor_(XLATensor& input);
 
+  static XLATensor trunc(const XLATensor& input);
+  static void trunc_(XLATensor& input);
+
   static XLATensor slice(const XLATensor& input, xla::int64 dim,
                          xla::int64 start, xla::int64 end, xla::int64 step);