Promote sides of bitwise logical operations

Author: asuhan

test/cpp/test_aten_xla_tensor.cpp

@@ -5514,6 +5514,19 @@ TEST_F(AtenXlaTensorTest, TestBitwiseAndScalarInPlace) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestBitwiseAndPromotion) {
+  at::Tensor input = at::rand({4, 2}, at::TensorOptions(at::kFloat));
+  at::Tensor view = input.reshape(-1);
+  at::Tensor result = at::__and__(view.gt(0), view.ne(0));
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_input = torch::autograd::make_variable(
+        bridge::CreateXlaTensor(input, device), false);
+    at::Tensor xla_view = xla_input.reshape(-1);
+    at::Tensor xla_result = at::__and__(xla_view.gt(0), xla_view.ne(0));
+    EXPECT_TRUE(EqualValues(result, xla_result));
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestBitwiseOr) {
   at::Tensor lhs = at::randint(0, std::numeric_limits<int32_t>::max(), {4, 2},
                                at::TensorOptions(at::kInt));

torch_xla/csrc/ops/bitwise_ir_ops.cpp

@@ -12,7 +12,10 @@ Value BitwiseAnd(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));
-    return node.ReturnOp(op0 & op1, loctx);
+    auto kernel = [](const xla::XlaOp& op0, const xla::XlaOp& op1) {
+      return op0 & op1;
+    };
+    return node.ReturnOp(XlaHelpers::PromotedBinaryOp(op0, op1, kernel), loctx);
   };
   return GenericOp(OpKind(at::aten::__and__), OpList{node1, node2},
                    node1.shape(), std::move(lower_fn));
@@ -22,7 +25,10 @@ Value BitwiseOr(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));
-    return node.ReturnOp(op0 | op1, loctx);
+    auto kernel = [](const xla::XlaOp& op0, const xla::XlaOp& op1) {
+      return op0 | op1;
+    };
+    return node.ReturnOp(XlaHelpers::PromotedBinaryOp(op0, op1, kernel), loctx);
   };
   return GenericOp(OpKind(at::aten::__or__), OpList{node1, node2},
                    node1.shape(), std::move(lower_fn));
@@ -32,7 +38,10 @@ Value BitwiseXor(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));
-    return node.ReturnOp(op0 ^ op1, loctx);
+    auto kernel = [](const xla::XlaOp& op0, const xla::XlaOp& op1) {
+      return op0 ^ op1;
+    };
+    return node.ReturnOp(XlaHelpers::PromotedBinaryOp(op0, op1, kernel), loctx);
   };
   return GenericOp(OpKind(at::aten::__xor__), OpList{node1, node2},
                    node1.shape(), std::move(lower_fn));