WIP: Add bitwise left/right

kernels/portable/cpu/op_bitwise_left_shift.cpp

@@ -0,0 +1,38 @@
+/*
+ * 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_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.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,37 @@
+/*
+ * 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_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.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

@@ -26,6 +26,23 @@ namespace internal {
 DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_and, &)
 DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_or, |)
 DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_xor, ^)
+DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_left_shift, <<)
+DEFINE_BINARY_OPERATOR_TEMPLATE(bitwise_right_shift, >>)
+
+// 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 lhs << rhs;
+  }
+};
+
+template <typename T = void>
+struct bit_rshift {
+  constexpr T operator()(const T& lhs, const T& rhs) const {
+    return lhs >> rhs;
+  }
+};
 
 template <typename T>
 using bitwise_fn = T (*)(const T, const T);
@@ -42,6 +59,12 @@ constexpr bitwise_fn<T> get_bitwise_fn() {
   if (op == "bitwise_xor.Tensor_out" || op == "bitwise_xor.Scalar_out") {
     return bitwise_xor;
   }
+  if (op == "bitwise_left_shift.Tensor_out" || op == "bitwise_left_shift.Scalar_out") {
+    return bitwise_left_shift;
+  }
+  if (op == "bitwise_right_shift.Tensor_out" || op == "bitwise_right_shift.Scalar_out") {
+    return bitwise_right_shift;
+  }
   return nullptr;
 };
 

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.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_left_shift_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.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::bitwise_right_shift_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/op_bitwise_left_shift_test.cpp

@@ -0,0 +1,119 @@
+/*
+ * 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}));
+}

kernels/test/op_bitwise_right_shift_test.cpp

@@ -0,0 +1,119 @@
+/*
+ * 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 OpBitwiseRightShiftTensorOutTest : public OperatorTest {
+ protected:
+  Tensor& op_bitwise_right_shift_tensor_out(
+      const Tensor& self,
+      const Tensor& other,
+      Tensor& out) {
+    return torch::executor::aten::bitwise_right_shift_outf(context_, self, other, out);
+  }
+};
+
+class OpBitwiseRightShiftScalarOutTest : public OperatorTest {
+ protected:
+  Tensor& op_bitwise_right_shift_scalar_out(
+      const Tensor& self,
+      const Scalar& other,
+      Tensor& out) {
+    return torch::executor::aten::bitwise_right_shift_outf(context_, self, other, out);
+  }
+};
+
+TEST_F(OpBitwiseRightShiftTensorOutTest, SmokeTestInt) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Test basic right shift: [8, 16, 32, 64] >> [1, 2, 1, 3] = [4, 4, 16, 8]
+  Tensor a = tf.make({2, 2}, {8, 16, 32, 64});
+  Tensor b = tf.make({2, 2}, {1, 2, 1, 3});
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_right_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {4, 4, 16, 8}));
+}
+
+TEST_F(OpBitwiseRightShiftTensorOutTest, SmokeTestByte) {
+  TensorFactory<ScalarType::Byte> tf;
+
+  // Test with byte values: [128, 64, 32, 16] >> [1, 1, 2, 3] = [64, 32, 8, 2]
+  Tensor a = tf.make({2, 2}, {128, 64, 32, 16});
+  Tensor b = tf.make({2, 2}, {1, 1, 2, 3});
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_right_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {64, 32, 8, 2}));
+}
+
+TEST_F(OpBitwiseRightShiftTensorOutTest, 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_right_shift_tensor_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {5, 10, 15, 20}));
+}
+
+TEST_F(OpBitwiseRightShiftScalarOutTest, SmokeTestInt) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Test shifting by scalar: [16, 32, 48, 64] >> 2 = [4, 8, 12, 16]
+  Tensor a = tf.make({2, 2}, {16, 32, 48, 64});
+  Scalar b = 2;
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_right_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {4, 8, 12, 16}));
+}
+
+TEST_F(OpBitwiseRightShiftScalarOutTest, ShiftByOne) {
+  TensorFactory<ScalarType::Int> tf;
+
+  // Shifting by 1 should halve the value (integer division)
+  Tensor a = tf.make({2, 2}, {100, 50, 20, 10});
+  Scalar b = 1;
+
+  Tensor out = tf.zeros({2, 2});
+
+  op_bitwise_right_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {50, 25, 10, 5}));
+}
+
+TEST_F(OpBitwiseRightShiftScalarOutTest, 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_right_shift_scalar_out(a, b, out);
+  EXPECT_TENSOR_EQ(out, tf.make({2, 2}, {7, 14, 21, 28}));
+}