Add lowering for bitwise_left_shift. (#8865)

Author: ysiraichi

codegen/xla_native_functions.yaml

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

docs/source/contribute/codegen_migration.md

@@ -330,6 +330,7 @@ being incremented correctly.
 
 ## Sample PRs
 
+-   Lowering of `bitwise_left_shift` <https://github.com/pytorch/xla/pull/8865>
 -   Unary/Binary OP -\> Codegen erf, erfc, erfinv, and exp
     (<https://github.com/pytorch/xla/pull/3659>)
 -   OP with optional -\> Codegen binary_cross_entropy/backward

test/test_operations.py

@@ -35,6 +35,7 @@
     all_types_and_complex_and,
     all_types_and,
 )
+import torch.utils._pytree as pytree
 import torch_xla
 import torch_xla.core.xla_builder as xb
 import torch_xla.core.xla_op_registry as xor
@@ -2384,24 +2385,32 @@ 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)
-
+  def _test_no_fallback(self, runf, args):
     met.clear_all()
 
     def run(device):
-      return torch.conj(tensor.to(device))
+      args_ = pytree.tree_map_only(torch.Tensor,
+                                   lambda t: t.clone().detach().to(device),
+                                   args)
+      return runf(*args_)
 
     actual = run("cpu")
     expected = run(xm.xla_device())
 
-    self.assertEqual(
-        met.executed_fallback_ops(), [],
-        message="expected no fallback operations.")
+    self.assertFalse(
+        met.executed_fallback_ops(), msg="expected no fallback operations.")
     self.assertEqual(
         actual, expected.cpu(), message="XLA results should match CPU results.")
 
+  def test_conj_no_fallback(self):
+    tensor = torch.rand(2, 2, dtype=torch.complex64)
+    self._test_no_fallback(torch.conj, (tensor,))
+
+  def test_bitwise_left_shift_no_fallback(self):
+    t1 = torch.randint(0, 10, (2, 2))
+    t2 = torch.randint(0, 10, (2,))
+    self._test_no_fallback(torch.bitwise_left_shift, (t1, t2))
+
 
 class MNISTComparator(nn.Module):
 

torch_xla/csrc/ops/ops_lower_fn.cpp

@@ -277,6 +277,20 @@ torch_xla::XlaOpVector BitwiseXorTensor::Lower(LoweringContext* loctx) const {
                   loctx);
 }
 
+torch_xla::XlaOpVector BitwiseLeftShiftTensor::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) {
+                        return xla::ShiftLeft(
+                            one, two,
+                            XlaHelpers::getBroadcastDimensions(one, two));
+                      }),
+                  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

@@ -360,6 +360,14 @@ xla::Shape BitwiseXorTensorOutputShape(const torch::lazy::Value& input,
   });
 }
 
+xla::Shape BitwiseLeftShiftTensorOutputShape(const torch::lazy::Value& input,
+                                             const torch::lazy::Value& other) {
+  return InferBinaryOpShape(input, other, [](xla::XlaOp one, xla::XlaOp two) {
+    return xla::ShiftLeft(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

@@ -93,6 +93,9 @@ xla::Shape BitwiseOrTensorOutputShape(const torch::lazy::Value& input,
 xla::Shape BitwiseXorTensorOutputShape(const torch::lazy::Value& input,
                                        const torch::lazy::Value& other);
 
+xla::Shape BitwiseLeftShiftTensorOutputShape(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,