【Hackathon 5th No.34】为 Paddle 新增 bitwise_right_shift / bitwise_right_shift_ / bitwise_left_shift / bitwise_left_shift_ API -part

paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc

@@ -33,10 +33,21 @@ void OperatorDialect::initialize() {
   // NOTE(chenxi67): GET_OP_LIST is defined in cinn_op.h which is
   // generated by op_gen.py, see details in
   // paddle/cinn/hlir/dialect/CMakeLists.txt.
+#ifdef WIN32
+  RegisterOps<
+#define GET_OP_LIST1
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
+      >();
+  RegisterOps<
+#define GET_OP_LIST2
+#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
+      >();
+#else
   RegisterOps<
 #define GET_OP_LIST
 #include "paddle/cinn/hlir/dialect/operator/ir/cinn_op_info.cc"  // NOLINT
       >();
+#endif
   RegisterOp<GroupOp>();
   RegisterOp<ConcatOp>();
   RegisterOp<SplitOp>();

paddle/fluid/pir/dialect/op_generator/op_gen.py

@@ -168,9 +168,20 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 {define_type_id}
 """
 
-CC_OP_INFO_FILE_TEMPLATE = """#ifdef GET_OP_LIST
+CC_OP_INFO_FILE_TEMPLATE_PART1 = """#ifdef GET_OP_LIST
 #undef GET_OP_LIST
 {op_declare}
+"""
+
+CC_OP_INFO_FILE_TEMPLATE_WIN_PART1 = """#ifdef GET_OP_LIST1
+#undef GET_OP_LIST1
+{op_declare_first_part}
+#elif defined(GET_OP_LIST2)
+#undef GET_OP_LIST2
+{op_declare_second_part}
+"""
+
+CC_OP_INFO_FILE_TEMPLATE_PART2 = """
 #else
 // This file is generated by "paddle/fluid/pir/dialect/op_generator/op_gen.py"
 #include "{h_file}"
@@ -1994,11 +2005,33 @@ def OpGenerator(
         op_to_multi_kernels_map_str = ""
 
     if op_info_file is not None:
-        op_info_str = CC_OP_INFO_FILE_TEMPLATE.format(
-            op_declare=",".join(op_list_strs).replace("\n", ""),
-            op_to_multi_kernels_map=op_to_multi_kernels_map_str,
-            h_file=op_def_h_file[:-4],
-        )
+        if sys.platform == "win32":
+            n = len(op_list_strs) // 2
+            first_part_op_info = op_list_strs[:n]
+            second_part_op_info = op_list_strs[n:]
+            CC_OP_INFO_FILE_TEMPLATE = (
+                CC_OP_INFO_FILE_TEMPLATE_WIN_PART1
+                + CC_OP_INFO_FILE_TEMPLATE_PART2
+            )
+            op_info_str = CC_OP_INFO_FILE_TEMPLATE.format(
+                op_declare_first_part=",".join(first_part_op_info).replace(
+                    "\n", ""
+                ),
+                op_declare_second_part=",".join(second_part_op_info).replace(
+                    "\n", ""
+                ),
+                op_to_multi_kernels_map=op_to_multi_kernels_map_str,
+                h_file=op_def_h_file[:-4],
+            )
+        else:
+            CC_OP_INFO_FILE_TEMPLATE = (
+                CC_OP_INFO_FILE_TEMPLATE_PART1 + CC_OP_INFO_FILE_TEMPLATE_PART2
+            )
+            op_info_str = CC_OP_INFO_FILE_TEMPLATE.format(
+                op_declare=",".join(op_list_strs).replace("\n", ""),
+                op_to_multi_kernels_map=op_to_multi_kernels_map_str,
+                h_file=op_def_h_file[:-4],
+            )
 
         with open(op_info_file, 'w') as f:
             f.write(op_info_str)

paddle/fluid/pir/dialect/operator/ir/op_dialect.cc

@@ -84,11 +84,22 @@ void OperatorDialect::initialize() {
   // paddle/fluid/pir/dialect/CMakeLists.txt.
   // NOTE(Ruting)GET_MANUAL_OP_LIST is define in manual_op.h"
   // use RegisterOps when list has more than two ops.
+#ifdef WIN32
   RegisterOps<
-#define GET_OP_LIST
+#define GET_OP_LIST1
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
       >();
 
+  RegisterOps<
+#define GET_OP_LIST2
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
+      >();
+#else
+  RegisterOps<
+#define GET_OP_LIST
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op_info.cc"  // NOLINT
+      >();
+#endif
   RegisterOps<
 #define GET_OP_LIST
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc"  // NOLINT

paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.cc

@@ -44,10 +44,21 @@ void OneDNNOperatorDialect::initialize() {
   // paddle/fluid/pir/dialect/CMakeLists.txt.
   // NOTE(Ruting)GET_MANUAL_OP_LIST is define in manual_op.h"
   // use RegisterOps when list has more than two ops.
+#ifdef WIN32
+  RegisterOps<
+#define GET_OP_LIST1
+#include "paddle/fluid/pir/dialect/operator/ir/pd_onednn_op_info.cc"  // NOLINT
+      >();
+  RegisterOps<
+#define GET_OP_LIST2
+#include "paddle/fluid/pir/dialect/operator/ir/pd_onednn_op_info.cc"  // NOLINT
+      >();
+#else
   RegisterOps<
 #define GET_OP_LIST
 #include "paddle/fluid/pir/dialect/operator/ir/pd_onednn_op_info.cc"  // NOLINT
       >();
+#endif
 }
 
 void OneDNNOperatorDialect::PrintType(pir::Type type, std::ostream &os) const {

paddle/phi/api/yaml/ops.yaml

@@ -343,6 +343,26 @@
     backend : x
   inplace: (x -> out)
 
+- op : bitwise_left_shift_arithmetic
+  args : (Tensor x, Tensor y)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+  kernel :
+    func : bitwise_left_shift_arithmetic
+    backend : x
+  inplace: (x -> out)
+
+- op : bitwise_left_shift_logic
+  args : (Tensor x, Tensor y)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+  kernel :
+    func : bitwise_left_shift_logic
+    backend : x
+  inplace: (x -> out)
+
 - op : bitwise_not
   args : (Tensor x)
   output : Tensor(out)
@@ -364,6 +384,26 @@
     backend : x
   inplace: (x -> out)
 
+- op : bitwise_right_shift_arithmetic
+  args : (Tensor x, Tensor y)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+  kernel :
+    func : bitwise_right_shift_arithmetic
+    backend : x
+  inplace: (x -> out)
+
+- op : bitwise_right_shift_logic
+  args : (Tensor x, Tensor y)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+  kernel :
+    func : bitwise_right_shift_logic
+    backend : x
+  inplace: (x -> out)
+
 - op : bitwise_xor
   args : (Tensor x, Tensor y)
   output : Tensor(out)

paddle/phi/kernels/bitwise_kernel.h

@@ -41,4 +41,28 @@ void BitwiseNotKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       DenseTensor* out);
 
+template <typename T, typename Context>
+void BitwiseLeftShiftArithmeticKernel(const Context& dev_ctx,
+                                      const DenseTensor& x,
+                                      const DenseTensor& y,
+                                      DenseTensor* out);
+
+template <typename T, typename Context>
+void BitwiseLeftShiftLogicKernel(const Context& dev_ctx,
+                                 const DenseTensor& x,
+                                 const DenseTensor& y,
+                                 DenseTensor* out);
+
+template <typename T, typename Context>
+void BitwiseRightShiftArithmeticKernel(const Context& dev_ctx,
+                                       const DenseTensor& x,
+                                       const DenseTensor& y,
+                                       DenseTensor* out);
+
+template <typename T, typename Context>
+void BitwiseRightShiftLogicKernel(const Context& dev_ctx,
+                                  const DenseTensor& x,
+                                  const DenseTensor& y,
+                                  DenseTensor* out);
+
 }  // namespace phi

paddle/phi/kernels/cpu/bitwise_kernel.cc

@@ -40,6 +40,31 @@ DEFINE_BITWISE_KERNEL(Or)
 DEFINE_BITWISE_KERNEL(Xor)
 #undef DEFINE_BITWISE_KERNEL
 
+#define DEFINE_BITWISE_KERNEL_WITH_INVERSE(op_type)                         \
+  template <typename T, typename Context>                                   \
+  void Bitwise##op_type##Kernel(const Context& dev_ctx,                     \
+                                const DenseTensor& x,                       \
+                                const DenseTensor& y,                       \
+                                DenseTensor* out) {                         \
+    funcs::Bitwise##op_type##Functor<T> func;                               \
+    funcs::InverseBitwise##op_type##Functor<T> inv_func;                    \
+    auto x_dims = x.dims();                                                 \
+    auto y_dims = y.dims();                                                 \
+    if (x_dims.size() >= y_dims.size()) {                                   \
+      funcs::ElementwiseCompute<funcs::Bitwise##op_type##Functor<T>, T>(    \
+          dev_ctx, x, y, func, out);                                        \
+    } else {                                                                \
+      funcs::ElementwiseCompute<funcs::InverseBitwise##op_type##Functor<T>, \
+                                T>(dev_ctx, x, y, inv_func, out);           \
+    }                                                                       \
+  }
+
+DEFINE_BITWISE_KERNEL_WITH_INVERSE(LeftShiftArithmetic)
+DEFINE_BITWISE_KERNEL_WITH_INVERSE(LeftShiftLogic)
+DEFINE_BITWISE_KERNEL_WITH_INVERSE(RightShiftArithmetic)
+DEFINE_BITWISE_KERNEL_WITH_INVERSE(RightShiftLogic)
+#undef DEFINE_BITWISE_KERNEL_WITH_INVERSE
+
 template <typename T, typename Context>
 void BitwiseNotKernel(const Context& dev_ctx,
                       const DenseTensor& x,
@@ -97,3 +122,43 @@ PD_REGISTER_KERNEL(bitwise_not,
                    int16_t,
                    int,
                    int64_t) {}
+
+PD_REGISTER_KERNEL(bitwise_left_shift_arithmetic,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BitwiseLeftShiftArithmeticKernel,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(bitwise_left_shift_logic,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BitwiseLeftShiftLogicKernel,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(bitwise_right_shift_arithmetic,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BitwiseRightShiftArithmeticKernel,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(bitwise_right_shift_logic,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BitwiseRightShiftLogicKernel,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}