Integer division binary ops: floor_divide, fmod, remainder (#6012)

Summary:
Pull Request resolved: #6012

- remainder:      386 K -> 15 K
- fmod:               317 K -> 14 K
- floor_divide:    255 K -> 11 K
ghstack-source-id: 246985124
exported-using-ghexport

Reviewed By: swolchok

Differential Revision: D63909727

fbshipit-source-id: 44d1fd575a3723df5982833abf5f9528b307a4bf

Author: manuelcandales

kernels/portable/cpu/op_floor_divide.cpp

@@ -6,7 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
+#include <executorch/kernels/portable/cpu/util/elementwise_util.h>
 #include <executorch/kernels/portable/cpu/util/math_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <executorch/runtime/platform/assert.h>
@@ -17,106 +17,61 @@ namespace torch {
 namespace executor {
 namespace native {
 
-using Tensor = exec_aten::Tensor;
-using ScalarType = exec_aten::ScalarType;
-
-namespace {
-template <
-    bool can_cast,
-    typename CTYPE_A,
-    typename CTYPE_B,
-    typename CTYPE_IN,
-    typename CTYPE_OUT>
-struct FloorDivideInner;
-
-template <
-    typename CTYPE_A,
-    typename CTYPE_B,
-    typename CTYPE_IN,
-    typename CTYPE_OUT>
-struct FloorDivideInner<true, CTYPE_A, CTYPE_B, CTYPE_IN, CTYPE_OUT> {
-  static void
-  run(const Tensor& a, const Tensor& b, Tensor& out, bool& div_by_zero_error) {
-    apply_binary_elementwise_fn<CTYPE_A, CTYPE_B, CTYPE_OUT>(
-        // NOLINTNEXTLINE(facebook-hte-ConstantArgumentPassByValue)
-        [&div_by_zero_error](const CTYPE_A val_a, const CTYPE_B val_b) {
-          if (is_integral_type<CTYPE_IN, /*includeBool=*/true>::value) {
-            if (val_b == 0) {
-              div_by_zero_error = true;
-              return static_cast<CTYPE_OUT>(0);
-            }
-          }
-          CTYPE_IN a_casted = static_cast<CTYPE_IN>(val_a);
-          CTYPE_IN b_casted = static_cast<CTYPE_IN>(val_b);
-          CTYPE_IN value = utils::floor_divide<CTYPE_IN>(a_casted, b_casted);
-
-          return static_cast<CTYPE_OUT>(value);
-        },
-        a,
-        b,
-        out);
-  }
-};
-
-struct ReportCanCastBug {
-  static void run(const Tensor&, const Tensor&, Tensor&, bool&) {
-    ET_DCHECK_MSG(false, "BUG: canCast should have been checked above");
-  }
-};
-
-template <
-    typename CTYPE_A,
-    typename CTYPE_B,
-    typename CTYPE_IN,
-    typename CTYPE_OUT>
-struct FloorDivideInner<false, CTYPE_A, CTYPE_B, CTYPE_IN, CTYPE_OUT>
-    : public ReportCanCastBug {};
-
-} // namespace
-
 Tensor& floor_divide_out(
     KernelRuntimeContext& ctx,
     const Tensor& a,
     const Tensor& b,
     Tensor& out) {
+  // Common Dtype
+  ScalarType common_type = promoteTypes(a.scalar_type(), b.scalar_type());
+
+  // Check Common Dtype
   ET_KERNEL_CHECK(
       ctx,
-      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
+      (canCast(common_type, out.scalar_type()) &&
+       common_type != ScalarType::Bool),
       InvalidArgument,
       out);
 
-  ET_KERNEL_CHECK(ctx, tensor_is_real_type(out), InvalidArgument, out);
-
+  // Check Dim Order
   ET_KERNEL_CHECK(
       ctx, tensors_have_same_dim_order(a, b, out), InvalidArgument, out);
 
-  ScalarType a_type = a.scalar_type();
-  ScalarType b_type = b.scalar_type();
-  ScalarType common_type = promoteTypes(a_type, b_type);
-  ScalarType out_type = out.scalar_type();
+  // Resize
+  ET_KERNEL_CHECK(
+      ctx,
+      resize_to_broadcast_target_size(a, b, out) == Error::Ok,
+      InvalidArgument,
+      out);
+
+  // Compute Dtype
+  ScalarType compute_type = utils::get_compute_type(common_type);
 
-  ET_KERNEL_CHECK(ctx, canCast(common_type, out_type), InvalidArgument, out);
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "floor_divide.out";
 
-  auto div_by_zero_error = false;
+  bool div_by_zero_error = false;
 
-  ET_SWITCH_REAL_TYPES_AND(
-      Bool, a_type, ctx, "floor_divide.out", CTYPE_A, [&]() {
-        ET_SWITCH_REAL_TYPES_AND(
-            Bool, b_type, ctx, "floor_divide.out", CTYPE_B, [&]() {
-              using CTYPE_IN = typename torch::executor::
-                  promote_types<CTYPE_A, CTYPE_B>::type;
-              ET_DCHECK(CppTypeToScalarType<CTYPE_IN>::value == common_type);
-              ET_SWITCH_REAL_TYPES(
-                  out_type, ctx, "floor_divide.out", CTYPE_OUT, [&]() {
-                    FloorDivideInner<
-                        can_cast<CTYPE_IN, CTYPE_OUT>::value,
-                        CTYPE_A,
-                        CTYPE_B,
-                        CTYPE_IN,
-                        CTYPE_OUT>::run(a, b, out, div_by_zero_error);
-                  });
-            });
-      });
+  ET_SWITCH_REAL_TYPES(compute_type, ctx, op_name, CTYPE_COMPUTE, [&]() {
+    utils::apply_bitensor_elementwise_fn<CTYPE_COMPUTE, op_name>(
+        [&div_by_zero_error](
+            const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
+          if (is_integral_type<CTYPE_COMPUTE, /*includeBool=*/true>::value) {
+            if (val_b == 0) {
+              div_by_zero_error = true;
+              return static_cast<CTYPE_COMPUTE>(0);
+            }
+          }
+          return utils::floor_divide(val_a, val_b);
+        },
+        ctx,
+        a,
+        utils::SupportedTensorDtypes::REALHBBF16,
+        b,
+        utils::SupportedTensorDtypes::REALHBBF16,
+        out,
+        utils::SupportedTensorDtypes::REALHBF16);
+  });
 
   ET_KERNEL_CHECK_MSG(
       ctx,