Add lowering for bitwise_right_shift. (#8866)

Author: ysiraichi

codegen/xla_native_functions.yaml

@@ -30,6 +30,7 @@ full_codegen:
   - binary_cross_entropy_backward
   - bitwise_not
   - bitwise_left_shift.Tensor
+  - bitwise_right_shift.Tensor
   - ceil
   - cholesky
   - clamp.Tensor

test/test_operations.py

@@ -2411,6 +2411,11 @@ def test_bitwise_left_shift_no_fallback(self):
     t2 = torch.randint(0, 10, (2,))
     self._test_no_fallback(torch.bitwise_left_shift, (t1, t2))
 
+  def test_bitwise_right_shift_no_fallback(self):
+    t1 = torch.randint(0, 10, (2, 2))
+    t2 = torch.randint(0, 10, (2,))
+    self._test_no_fallback(torch.bitwise_right_shift, (t1, t2))
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -291,6 +291,25 @@ torch_xla::XlaOpVector BitwiseLeftShiftTensor::Lower(
                   loctx);
 }
 
+torch_xla::XlaOpVector BitwiseRightShiftTensor::Lower(
+    LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_other_input = loctx->GetOutputOp(operand(1));
+  return ReturnOp(
+      XlaHelpers::PromotedBinaryOp(
+          xla_input, xla_other_input,
+          [](xla::XlaOp one, xla::XlaOp two) {
+            auto broadcast_dims = XlaHelpers::getBroadcastDimensions(one, two);
+            if (xla::primitive_util::IsSignedIntegralType(
+                    XlaHelpers::TypeOfXlaOp(one))) {
+              return xla::ShiftRightArithmetic(one, two, broadcast_dims);
+            } else {
+              return xla::ShiftRightLogical(one, two, broadcast_dims);
+            }
+          }),
+      loctx);
+}
+
 torch_xla::XlaOpVector Ceil::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

@@ -368,6 +368,14 @@ xla::Shape BitwiseLeftShiftTensorOutputShape(const torch::lazy::Value& input,
   });
 }
 
+xla::Shape BitwiseRightShiftTensorOutputShape(const torch::lazy::Value& input,
+                                              const torch::lazy::Value& other) {
+  return InferBinaryOpShape(input, other, [](xla::XlaOp one, xla::XlaOp two) {
+    return xla::ShiftRightArithmetic(
+        one, two, XlaHelpers::getBroadcastDimensions(one, two));
+  });
+}
+
 xla::Shape CeilOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -96,6 +96,9 @@ xla::Shape BitwiseXorTensorOutputShape(const torch::lazy::Value& input,
 xla::Shape BitwiseLeftShiftTensorOutputShape(const torch::lazy::Value& input,
                                              const torch::lazy::Value& other);
 
+xla::Shape BitwiseRightShiftTensorOutputShape(const torch::lazy::Value& input,
+                                              const torch::lazy::Value& other);
+
 xla::Shape CeilOutputShape(const torch::lazy::Value& input);
 
 xla::Shape CholeskyOutputShape(const torch::lazy::Value& input,