floor_divide fixed for signed 0 output

Author: ndgrigorian

dpctl/tensor/libtensor/include/kernels/elementwise_functions/floor_divide.hpp

@@ -60,24 +60,25 @@ struct FloorDivideFunctor
 
     resT operator()(const argT1 &in1, const argT2 &in2)
     {
-        auto tmp = in1 / in2;
-        if constexpr (std::is_integral_v<decltype(tmp)>) {
-            if constexpr (std::is_unsigned_v<decltype(tmp)>) {
-                return (in2 == argT2(0)) ? resT(0) : tmp;
+        if constexpr (std::is_integral_v<argT1> || std::is_integral_v<argT2>) {
+            static_assert(std::is_same_v<argT1, argT2>);
+            if (in2 == 0) {
+                return resT(0);
+            }
+            auto tmp = in1 / in2;
+            if constexpr (std::is_unsigned_v<argT1> ||
+                          std::is_unsigned_v<argT2>) {
+                return tmp;
             }
             else {
-                if (in2 == argT2(0)) {
-                    return resT(0);
-                }
-                else {
-                    auto rem = in1 % in2;
-                    auto corr = (rem != 0 && ((rem < 0) != (in2 < 0)));
-                    return (tmp - corr);
-                }
+                auto rem = in1 % in2;
+                auto corr = (rem != 0 && ((rem < 0) != (in2 < 0)));
+                return (tmp - corr);
             }
         }
         else {
-            return sycl::floor(tmp);
+            auto tmp = in1 / in2;
+            return (tmp == 0) ? resT(tmp) : resT(std::floor(tmp));
         }
     }
 
@@ -88,26 +89,26 @@ struct FloorDivideFunctor
         auto tmp = in1 / in2;
         using tmpT = typename decltype(tmp)::element_type;
         if constexpr (std::is_integral_v<tmpT>) {
-            if constexpr (std::is_signed_v<tmpT>) {
-                auto rem_tmp = in1 % in2;
+            if constexpr (std::is_unsigned_v<tmpT>) {
 #pragma unroll
                 for (int i = 0; i < vec_sz; ++i) {
                     if (in2[i] == argT2(0)) {
                         tmp[i] = tmpT(0);
                     }
-                    else {
-                        tmpT corr = (rem_tmp[i] != 0 &&
-                                     ((rem_tmp[i] < 0) != (in2[i] < 0)));
-                        tmp[i] -= corr;
-                    }
                 }
             }
             else {
+                auto rem = in1 % in2;
 #pragma unroll
                 for (int i = 0; i < vec_sz; ++i) {
-                    if (in2[i] == argT2(0)) {
+                    if (in2[i] == 0) {
                         tmp[i] = tmpT(0);
                     }
+                    else {
+                        tmpT corr =
+                            (rem[i] != 0 && ((rem[i] < 0) != (in2[i] < 0)));
+                        tmp[i] -= corr;
+                    }
                 }
             }
             if constexpr (std::is_same_v<resT, tmpT>) {
@@ -119,16 +120,18 @@ struct FloorDivideFunctor
             }
         }
         else {
-            sycl::vec<resT, vec_sz> res = sycl::floor(tmp);
-            if constexpr (std::is_same_v<resT,
-                                         typename decltype(res)::element_type>)
-            {
-                return res;
+#pragma unroll
+            for (int i = 0; i < vec_sz; ++i) {
+                if (in2[i] != 0) {
+                    tmp[i] = std::floor(tmp[i]);
+                }
+            }
+            if constexpr (std::is_same_v<resT, tmpT>) {
+                return tmp;
             }
             else {
                 using dpctl::tensor::type_utils::vec_cast;
-                return vec_cast<resT, typename decltype(res)::element_type,
-                                vec_sz>(res);
+                return vec_cast<resT, tmpT, vec_sz>(tmp);
             }
         }
     }

dpctl/tests/elementwise/test_floor_divide.py

@@ -203,16 +203,6 @@ def test_floor_divide_gh_1247():
         dpt.asnumpy(res), np.full(res.shape, -1, dtype=res.dtype)
     )
 
-    # attempt to invoke sycl::vec overload using a larger array
-    x = dpt.arange(-64, 65, 1, dtype="i4")
-    np.testing.assert_array_equal(
-        dpt.asnumpy(dpt.floor_divide(x, 3)), np.floor_divide(dpt.asnumpy(x), 3)
-    )
-    np.testing.assert_array_equal(
-        dpt.asnumpy(dpt.floor_divide(x, -3)),
-        np.floor_divide(dpt.asnumpy(x), -3),
-    )
-
 
 @pytest.mark.parametrize("dtype", _no_complex_dtypes[1:9])
 def test_floor_divide_integer_zero(dtype):
@@ -226,10 +216,42 @@ def test_floor_divide_integer_zero(dtype):
         dpt.asnumpy(res), np.zeros(x.shape, dtype=res.dtype)
     )
 
-    # attempt to invoke sycl::vec overload using a larger array
-    x = dpt.arange(129, dtype=dtype, sycl_queue=q)
-    y = dpt.zeros_like(x, sycl_queue=q)
+
+def test_floor_divide_special_cases():
+    q = get_queue_or_skip()
+
+    x = dpt.empty(1, dtype="f4", sycl_queue=q)
+    y = dpt.empty_like(x)
+    x[0], y[0] = dpt.inf, dpt.inf
+    res = dpt.floor_divide(x, y)
+    with np.errstate(all="ignore"):
+        res_np = np.floor_divide(dpt.asnumpy(x), dpt.asnumpy(y))
+        np.testing.assert_array_equal(dpt.asnumpy(res), res_np)
+
+    x[0], y[0] = 0.0, -1.0
+    res = dpt.floor_divide(x, y)
+    x_np = dpt.asnumpy(x)
+    y_np = dpt.asnumpy(y)
+    res_np = np.floor_divide(x_np, y_np)
+    np.testing.assert_array_equal(dpt.asnumpy(res), res_np)
+
+    res = dpt.floor_divide(y, x)
+    with np.errstate(all="ignore"):
+        res_np = np.floor_divide(y_np, x_np)
+        np.testing.assert_array_equal(dpt.asnumpy(res), res_np)
+
+    x[0], y[0] = -1.0, dpt.inf
     res = dpt.floor_divide(x, y)
     np.testing.assert_array_equal(
-        dpt.asnumpy(res), np.zeros(x.shape, dtype=res.dtype)
+        dpt.asnumpy(res), np.asarray([-0.0], dtype="f4")
     )
+
+    res = dpt.floor_divide(y, x)
+    np.testing.assert_array_equal(
+        dpt.asnumpy(res), np.asarray([-dpt.inf], dtype="f4")
+    )
+
+    x[0], y[0] = 1.0, dpt.nan
+    res = dpt.floor_divide(x, y)
+    res_np = np.floor_divide(dpt.asnumpy(x), dpt.asnumpy(y))
+    np.testing.assert_array_equal(dpt.asnumpy(res), res_np)