add array-array div logic

abhishek-iitmadras · abhishek-iitmadras · commit e8349985ef01 · 2025-07-06T17:59:52.000+05:30
diff --git a/numpy/_core/src/umath/loops_arithmetic.dispatch.cpp b/numpy/_core/src/umath/loops_arithmetic.dispatch.cpp
@@ -169,10 +169,107 @@ T floor_div(T n, T d) {
     }
     return r;
 }
+// General divide implementation for arrays
+template <typename T>
+void simd_divide_contig_signed(T* src1, T* src2, T* dst, npy_intp len) {
+    using D = hn::ScalableTag<T>;
+    const D d;
+    const size_t N = hn::Lanes(d);
+    bool raise_overflow = false;
+    bool raise_divbyzero = false;
+    const auto vec_zero = hn::Zero(d);
+    const auto vec_min_val = hn::Set(d, std::numeric_limits<T>::min());
+    const auto vec_neg_one = hn::Set(d, static_cast<T>(-1));
+    
+    size_t i = 0;
+    for (; i + N <= static_cast<size_t>(len); i += N) {
+        const auto vec_a = hn::LoadU(d, src1 + i);
+        const auto vec_b = hn::LoadU(d, src2 + i);
+        const auto b_is_zero = hn::Eq(vec_b, vec_zero);
+        const auto a_is_min = hn::Eq(vec_a, vec_min_val);
+        const auto b_is_neg_one = hn::Eq(vec_b, vec_neg_one);
+        const auto overflow_cond = hn::And(a_is_min, b_is_neg_one);
+        auto vec_div = hn::Div(vec_a, vec_b);
+        const auto vec_mul = hn::Mul(vec_div, vec_b);
+        const auto has_remainder = hn::Ne(vec_a, vec_mul);
+        const auto a_sign = hn::Lt(vec_a, vec_zero);
+        const auto b_sign = hn::Lt(vec_b, vec_zero);
+        const auto different_signs = hn::Xor(a_sign, b_sign);
+        auto adjustment = hn::And(different_signs, has_remainder);
+        vec_div = hn::IfThenElse(adjustment, 
+                                 hn::Sub(vec_div, hn::Set(d, static_cast<T>(1))), 
+                                 vec_div);
+        vec_div = hn::IfThenElse(b_is_zero, vec_zero, vec_div);
+        vec_div = hn::IfThenElse(overflow_cond, vec_min_val, vec_div);
+        hn::StoreU(vec_div, d, dst + i);
+        if (!raise_divbyzero && !hn::AllFalse(d, b_is_zero)) {
+            raise_divbyzero = true;
+        }
+        if (!raise_overflow && !hn::AllFalse(d, overflow_cond)) {
+            raise_overflow = true;
+        }
+    }
+    for (; i < static_cast<size_t>(len); i++) {
+        T a = src1[i];
+        T b = src2[i];
+        
+        if (b == 0) {
+            dst[i] = 0;
+            raise_divbyzero = true;
+        } 
+        else if (a == std::numeric_limits<T>::min() && b == -1) {
+            dst[i] = std::numeric_limits<T>::min();
+            raise_overflow = true;
+        } 
+        else {
+            dst[i] = floor_div(a, b);
+        }
+    }
+    
+    set_float_status(raise_overflow, raise_divbyzero);
+}
+// Unsigned division for arrays
+template <typename T>
+void simd_divide_contig_unsigned(T* src1, T* src2, T* dst, npy_intp len) {
+    using D = hn::ScalableTag<T>;
+    const D d;
+    const size_t N = hn::Lanes(d);
+    
+    bool raise_divbyzero = false;
+    const auto vec_zero = hn::Zero(d);
+    
+    size_t i = 0;
+    for (; i + N <= static_cast<size_t>(len); i += N) {
+        const auto vec_a = hn::LoadU(d, src1 + i);
+        const auto vec_b = hn::LoadU(d, src2 + i);
+        const auto b_is_zero = hn::Eq(vec_b, vec_zero);
+        auto vec_div = hn::Div(vec_a, vec_b);
+        vec_div = hn::IfThenElse(b_is_zero, vec_zero, vec_div);
+        hn::StoreU(vec_div, d, dst + i);
+        if (!raise_divbyzero && !hn::AllFalse(d, b_is_zero)) {
+            raise_divbyzero = true;
+        }
+    }
+    for (; i < static_cast<size_t>(len); i++) {
+        T a = src1[i];
+        T b = src2[i];
+        
+        if (b == 0) {
+            dst[i] = 0;
+            raise_divbyzero = true;
+        } else {
+            dst[i] = a / b;
+        }
+    }
+    
+    set_float_status(false, raise_divbyzero);
+}
+
 
 // Dispatch functions for signed integer division
 template <typename T>
 void TYPE_divide(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)) {
+    npy_clear_floatstatus();
     if (IS_BINARY_REDUCE) {
         BINARY_REDUCE_LOOP(T) {
             const T divisor = *reinterpret_cast<T*>(ip2);
@@ -189,8 +286,22 @@ void TYPE_divide(char **args, npy_intp const *dimensions, npy_intp const *steps,
         *reinterpret_cast<T*>(iop1) = io1;
         return;
     }
-#if NPY_SIMD   
-    if (IS_BLOCKABLE_BINARY_SCALAR2(sizeof(T), NPY_SIMD_WIDTH) &&
+#if NPY_SIMD
+    // Handle array-array case
+    if (IS_BLOCKABLE_BINARY(sizeof(T), NPY_SIMD_WIDTH)) 
+    {
+        bool no_overlap = nomemoverlap(args[2], steps[2], args[0], steps[0], dimensions[0]) &&
+                         nomemoverlap(args[2], steps[2], args[1], steps[1], dimensions[0]);
+        // Check if we can use SIMD for contiguous arrays - all steps must equal to sizeof(T)
+        if (steps[0] == sizeof(T) && steps[1] == sizeof(T) && steps[2] == sizeof(T) && no_overlap) {
+            T* src1 = (T*)args[0];
+            T* src2 = (T*)args[1];
+            T* dst = (T*)args[2];
+            simd_divide_contig_signed(src1, src2, dst, dimensions[0]);
+            return;
+        }
+    }
+    else if (IS_BLOCKABLE_BINARY_SCALAR2(sizeof(T), NPY_SIMD_WIDTH) &&
         *reinterpret_cast<T*>(args[1]) != 0)
     {
         bool no_overlap = nomemoverlap(args[2], steps[2], args[0], steps[0], dimensions[0]);
@@ -225,6 +336,7 @@ void TYPE_divide(char **args, npy_intp const *dimensions, npy_intp const *steps,
 // Dispatch functions for unsigned integer division
 template <typename T>
 void TYPE_divide_unsigned(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func)) {
+    npy_clear_floatstatus();
     if (IS_BINARY_REDUCE) {
         BINARY_REDUCE_LOOP(T) {
             const T d = *reinterpret_cast<T*>(ip2);
@@ -239,7 +351,20 @@ void TYPE_divide_unsigned(char **args, npy_intp const *dimensions, npy_intp cons
         return;
     }
 #if NPY_SIMD
-    if (IS_BLOCKABLE_BINARY_SCALAR2(sizeof(T), NPY_SIMD_WIDTH) &&
+    // Handle array-array case
+    if (IS_BLOCKABLE_BINARY(sizeof(T), NPY_SIMD_WIDTH)) {
+        bool no_overlap = nomemoverlap(args[2], steps[2], args[0], steps[0], dimensions[0]) &&
+                         nomemoverlap(args[2], steps[2], args[1], steps[1], dimensions[0]);
+        // Check if we can use SIMD for contiguous arrays - all steps must equal to sizeof(T)
+        if (steps[0] == sizeof(T) && steps[1] == sizeof(T) && steps[2] == sizeof(T) && no_overlap) {
+            T* src1 = (T*)args[0];
+            T* src2 = (T*)args[1];
+            T* dst = (T*)args[2];
+            simd_divide_contig_unsigned(src1, src2, dst, dimensions[0]);
+            return;
+        }
+    }
+    else if (IS_BLOCKABLE_BINARY_SCALAR2(sizeof(T), NPY_SIMD_WIDTH) &&
         *reinterpret_cast<T*>(args[1]) != 0)
     {
         bool no_overlap = nomemoverlap(args[2], steps[2], args[0], steps[0], dimensions[0]);
