Add bitwise left/right (pytorch#15893)

Fixes
pytorch#15673 (comment)

Did not add as core aten op. But used "func" and enabled exporting via:

```
to_edge_transform_and_lower(
        exported_program,
        partitioner=None,
        compile_config=EdgeCompileConfig(
            _core_aten_ops_exception_list=[
                torch.ops.aten.bitwise_left_shift.Tensor,
                torch.ops.aten.bitwise_left_shift.Tensor_Scalar,
                torch.ops.aten.bitwise_right_shift.Tensor,
                torch.ops.aten.bitwise_right_shift.Tensor_Scalar,
            ]
        ),
    )
```

In the future, we may promote to core aten op.

Author: mergennachin

kernels/portable/cpu/op_bitwise_left_shift.cpp

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/pattern/bitwise_op.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+Tensor& bitwise_left_shift_Tensor_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& a,
+    const Tensor& b,
+    Tensor& out) {
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "bitwise_left_shift.Tensor_out";
+  return internal::bitwise_tensor_out<internal::bit_lshift, op_name>(
+      ctx, a, b, out);
+}
+
+Tensor& bitwise_left_shift_Tensor_Scalar_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& a,
+    const Scalar& b,
+    Tensor& out) {
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] =
+      "bitwise_left_shift.Tensor_Scalar_out";
+  return internal::bitwise_scalar_out<internal::bit_lshift, op_name>(
+      ctx, a, b, out);
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/op_bitwise_right_shift.cpp

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/portable/cpu/pattern/bitwise_op.h>
+
+namespace torch {
+namespace executor {
+namespace native {
+
+Tensor& bitwise_right_shift_Tensor_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& a,
+    const Tensor& b,
+    Tensor& out) {
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "bitwise_right_shift.Tensor_out";
+  return internal::bitwise_tensor_out<internal::bit_rshift, op_name>(
+      ctx, a, b, out);
+}
+
+Tensor& bitwise_right_shift_Tensor_Scalar_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& a,
+    const Scalar& b,
+    Tensor& out) {
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] =
+      "bitwise_right_shift.Tensor_Scalar_out";
+  return internal::bitwise_scalar_out<internal::bit_rshift, op_name>(
+      ctx, a, b, out);
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch

kernels/portable/cpu/pattern/bitwise_op.h

@@ -27,6 +27,21 @@ DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_and, &)
 DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_or, |)
 DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_xor, ^)
 
+// Functor wrappers for shift operations (similar to std::bit_and, etc.)
+template <typename T = void>
+struct bit_lshift {
+  constexpr T operator()(const T& lhs, const T& rhs) const {
+    return static_cast<T>(lhs << rhs);
+  }
+};
+
+template <typename T = void>
+struct bit_rshift {
+  constexpr T operator()(const T& lhs, const T& rhs) const {
+    return static_cast<T>(lhs >> rhs);
+  }
+};
+
 template <typename T>
 using bitwise_fn = T (*)(const T, const T);
 

kernels/portable/functions.yaml

@@ -207,6 +207,26 @@
     - arg_meta: null
       kernel_name: torch::executor::bitwise_xor_Tensor_out
 
+- func: bitwise_left_shift.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_left_shift_Tensor_out
+
+- func: bitwise_left_shift.Tensor_Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_left_shift_Tensor_Scalar_out
+
+- func: bitwise_right_shift.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_right_shift_Tensor_out
+
+- func: bitwise_right_shift.Tensor_Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_right_shift_Tensor_Scalar_out
+
 - op: bmm.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -147,8 +147,10 @@ set(all_test_sources
     "op_atanh_test.cpp"
     "op_avg_pool2d_test.cpp"
     "op_bitwise_and_test.cpp"
+    "op_bitwise_left_shift_test.cpp"
     "op_bitwise_not_test.cpp"
     "op_bitwise_or_test.cpp"
+    "op_bitwise_right_shift_test.cpp"
     "op_bitwise_xor_test.cpp"
     "op_bmm_test.cpp"
     "op_cat_test.cpp"

kernels/test/TARGETS

@@ -9,6 +9,23 @@ oncall("executorch")
 
 define_common_targets()
 
+python_unittest(
+    name = "test_bitwise_shift",
+    srcs = ["test_bitwise_shift.py"],
+    preload_deps = [
+        "//executorch/kernels/portable:custom_ops_generated_lib",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:lib",
+        "//executorch/extension/export_util:export_util",
+        "//executorch/runtime:runtime",
+    ],
+    env = {
+        "PYTORCH_DISABLE_JUSTKNOBS": "1",
+    },
+)
+
 python_unittest(
     name = "gen_supported_features_test",
     srcs = ["gen_supported_features_test.py"],

kernels/test/op_bitwise_left_shift_test.cpp

@@ -0,0 +1,121 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
+#include <executorch/kernels/test/TestUtil.h>
+#include <executorch/runtime/core/exec_aten/exec_aten.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
+#include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
+
+#include <gtest/gtest.h>
+
+using namespace ::testing;
+using executorch::aten::Scalar;
+using executorch::aten::ScalarType;
+using executorch::aten::Tensor;
+using torch::executor::testing::TensorFactory;
+
+class OpBitwiseLeftShiftTensorOutTest : public OperatorTest {
+ protected:
+  Tensor& op_bitwise_left_shift_tensor_out(
+      const Tensor& self,
+      const Tensor& other,
+      Tensor& out) {
+    return torch::executor::aten::bitwise_left_shift_outf(
+        context_, self, other, out);
+  }
+};
+
+class OpBitwiseLeftShiftScalarOutTest : public OperatorTest {
+ protected:
+  Tensor& op_bitwise_left_shift_scalar_out(
+      const Tensor& self,
+      const Scalar& other,
+      Tensor& out) {
+    return torch::executor::aten::bitwise_left_shift_outf(
+        context_, self, other, out);
+  }
+};
+
+TEST_F(OpBitwiseLeftShiftTensorOutTest, SmokeTestInt) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Test basic left shift: [1, 2, 4, 8] << [1, 2, 1, 2] = [2, 8, 8, 32]
+  Tensor a = tf.make({2, 2}, {1, 2, 4, 8});
+  Tensor b = tf.make({2, 2}, {1, 2, 1, 2});
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {2, 8, 8, 32}));
+}
+
+TEST_F(OpBitwiseLeftShiftTensorOutTest, SmokeTestByte) {
+  TensorFactory<ScalarType::Byte> tf;
+
+  // Test with byte values: [1, 5, 10, 15] << [0, 1, 2, 3] = [1, 10, 40, 120]
+  Tensor a = tf.make({2, 2}, {1, 5, 10, 15});
+  Tensor b = tf.make({2, 2}, {0, 1, 2, 3});
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {1, 10, 40, 120}));
+}
+
+TEST_F(OpBitwiseLeftShiftTensorOutTest, ZeroShift) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Shifting by 0 should return the original value
+  Tensor a = tf.make({2, 2}, {5, 10, 15, 20});
+  Tensor b = tf.zeros({2, 2});
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {5, 10, 15, 20}));
+}
+
+TEST_F(OpBitwiseLeftShiftScalarOutTest, SmokeTestInt) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Test shifting by scalar: [1, 2, 3, 4] << 2 = [4, 8, 12, 16]
+  Tensor a = tf.make({2, 2}, {1, 2, 3, 4});
+  Scalar b = 2;
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {4, 8, 12, 16}));
+}
+
+TEST_F(OpBitwiseLeftShiftScalarOutTest, ShiftByOne) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Shifting by 1 should double the value
+  Tensor a = tf.make({2, 2}, {1, 5, 10, 100});
+  Scalar b = 1;
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {2, 10, 20, 200}));
+}
+
+TEST_F(OpBitwiseLeftShiftScalarOutTest, ShiftByZero) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Shifting by 0 should return the original value
+  Tensor a = tf.make({2, 2}, {7, 14, 21, 28});
+  Scalar b = 0;
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_left_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {7, 14, 21, 28}));
+}