Add ceil, trunc, sqrt, square,

numpy · Apr 28, 2024 · 7a9a4e9 · 7a9a4e9
1 parent 3d5580e
commit 7a9a4e9
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 8 deletions.
diff --git a/numpy/_core/src/umath/loop_unary_fp.dispatch.cpp b/numpy/_core/src/umath/loop_unary_fp.dispatch.cpp
@@ -9,7 +9,6 @@
 #include "numpy/npy_common.h"
 #include "numpy/npy_math.h"
 #include "numpy/utils.h"
-
 #include "fast_loop_macros.h"
 #include "loops_utils.h"
 #include <hwy/highway.h>
@@ -90,8 +89,14 @@ namespace hn = hwy::HWY_NAMESPACE;
     }                                                                          \
   }
 
+#define Square(x) hn::Mul(x, x) 
+//ceil, trunc, sqrt, square, 
 SUPER(Rint, hn::Round)
-
+SUPER(Floor, hn::Floor)
+SUPER(Ceil, hn::Ceil)
+SUPER(Trunc, hn::Trunc)
+SUPER(Sqrt, hn::Sqrt)
+SUPER(Square, Square)
 
 extern "C" {
 NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_rint)
@@ -105,4 +110,64 @@ NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_rint)
 {
   return SuperRint<npy_float>(args, dimensions, steps);
 }
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_floor)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperFloor<npy_double>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_floor)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperFloor<npy_float>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_ceil)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperCeil<npy_double>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_ceil)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperCeil<npy_float>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_trunc)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperTrunc<npy_double>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_trunc)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperTrunc<npy_float>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_sqrt)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperSqrt<npy_double>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_sqrt)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperSqrt<npy_float>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(DOUBLE_square)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperSquare<npy_double>(args, dimensions, steps);
+}
+
+NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(FLOAT_square)
+(char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))
+{
+  return SuperSquare<npy_float>(args, dimensions, steps);
+}
 }
diff --git a/numpy/_core/src/umath/loops_unary_fp.dispatch.c.src b/numpy/_core/src/umath/loops_unary_fp.dispatch.c.src
@@ -100,9 +100,9 @@ NPY_FINLINE double c_square_f64(double a)
  */
 #if @VCHK@
 /**begin repeat1
- * #kind     = floor, ceil, trunc, sqrt, absolute, square, reciprocal#
- * #intr     = floor, ceil, trunc, sqrt, abs,      square, recip#
- * #repl_0w1 = 0*6, 1#
+ * #kind     = absolute,reciprocal#
+ * #intr     = abs,recip#
+ * #repl_0w1 = 0, 1#
  */
 /**begin repeat2
  * #STYPE  = CONTIG, NCONTIG, CONTIG,  NCONTIG#
@@ -199,9 +199,9 @@ static void simd_@TYPE@_@kind@_@STYPE@_@DTYPE@
  * #VCHK = NPY_SIMD_F32, NPY_SIMD_F64#
  */
 /**begin repeat1
- * #kind  = floor, ceil, trunc, sqrt, absolute, square, reciprocal#
- * #intr  = floor, ceil, trunc, sqrt, abs,      square, recip#
- * #clear = 0,     0,    0,     0,    1,        0,      0#
+ * #kind  = absolute,reciprocal#
+ * #intr  = abs,recip#
+ * #clear = 1, 0#
  */
 NPY_NO_EXPORT void NPY_CPU_DISPATCH_CURFX(@TYPE@_@kind@)
 (char **args, npy_intp const *dimensions, npy_intp const *steps, void *NPY_UNUSED(func))