8000 ENH: add scalar special cases for boolean logical loops · numpy/numpy@d041977 · GitHub
[go: up one dir, main page]
Skip to content

Commit d041977

Browse files
juliantaylorjuliantaylor
committed
ENH: add scalar special cases for boolean logical loops
Scalar logical loops just boil down to memcpy or memset. While not very useful in general, some cases like masked arrays can profit when creating zero stride boolean arrays for readonly views.
1 parent f8d4a2e commit d041977

File tree

2 files changed

+86
-4
lines changed

2 files changed

+86
-4
lines changed

‎numpy/core/src/umath/loops.c.src

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -703,6 +703,27 @@ BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED
703703
}
704704
/**end repeat**/
705705

706+
#define BOOL_SCALAR_OP(cond, in_idx, inp, value) \
707+
if (!(cond)) { \
708+
if (steps[in_idx] == 1 && steps[2] == 1) { \
709+
memcpy(args[2], args[in_idx], dimensions[0]); \
710+
} \
711+
else { \
712+
BINARY_LOOP { \
713+
*op1 = *inp; \
714+
} \
715+
} \
716+
} \
717+
else { \
718+
if (steps[2] == 1) { \
719+
memset(args[2], value, dimensions[0]); \
720+
} \
721+
else { \
722+
BINARY_LOOP { \
723+
*op1 = value; \
724+
} \
725+
} \
726+
}
706727

707728
/**begin repeat
708729
* #kind = logical_and, logical_or#
@@ -763,7 +784,24 @@ BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED
763784
}
764785
}
765786
else {
766-
if (run_binary_simd_@kind@_BOOL(args, dimensions, steps)) {
787+
/*
788+
* scalar operations can occur e.g. in masked arrays where a
789+
* stride 0 array is used to represent no or full mask
790+
*/
791+
if (steps[0] == 0) {
792+
if (steps[1] == 0) {
793+
/* only the memset part of the loop */
794+
const npy_bool res = (*args[0] @OP@ *args[1]);
795+
BOOL_SCALAR_OP(1, 1, ip2, res);
796+
}
797+
else {
798+
BOOL_SCALAR_OP((*args[0] @SC@ 0), 1, ip2, !@and@);
799+
}
800+
}
801+
else if (steps[1] == 0) {
802+
BOOL_SCALAR_OP((*args[1] @SC@ 0), 0, ip1, !@and@);
803+
}
804+
else if (run_binary_simd_@kind@_BOOL(args, dimensions, steps)) {
767805
return;
768806
}
769807
else {
@@ -788,9 +826,14 @@ BOOL_@kind@(char **args, npy_intp *dimensions, npy_intp *steps, void *NPY_UNUSED
788826
return;
789827
}
790828
else {
791-
UNARY_LOOP {
792-
npy_bool in1 = *(npy_bool *)ip1;
793-
*((npy_bool *)op1) = in1 @OP@ 0;
829+
if (steps[0] == 0 && steps[1] == 1) {
830+
memset(args[1], *args[0] @OP@ 0, dimensions[0]);
831+
}
832+
else {
833+
UNARY_LOOP {
834+
npy_bool in1 = *(npy_bool *)ip1;
835+
*((npy_bool *)op1) = in1 @OP@ 0;
836+
}
794837
}
795838
}
796839
}

‎numpy/core/tests/test_numeric.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,17 @@ def test_logical_not_abs(self):
288288
np.abs(self.t, out=self.o)
289289
assert_array_equal(self.o, self.t)
290290

291+
# test scalar array specialization
292+
norm = np.zeros(114, dtype=bool)[1:]
293+
scalar_f = np.ndarray(buffer=np.zeros((), dtype=bool), strides=(0,),
294+
shape=norm.shape, dtype=bool)
295+
scalar_t = np.ndarray(buffer=np.ones((), dtype=bool), strides=(0,),
296+
shape=norm.shape, dtype=bool)
297+
assert_array_equal(~scalar_f, ~norm)
298+
assert_array_equal(np.abs(scalar_f), norm)
299+
assert_array_equal(~scalar_t, norm)
300+
assert_array_equal(np.abs(scalar_t), ~norm)
301+
291302
def test_logical_and_or_xor(self):
292303
assert_array_equal(self.t | self.t, self.t)
293304
assert_array_equal(self.f | self.f, self.f)
@@ -310,16 +321,44 @@ def test_logical_and_or_xor(self):
310321

311322
assert_array_equal(self.nm & self.t, self.nm)
312323
assert_array_equal(self.im & self.f, False)
324+
assert_array_equal(self.t & self.nm, self.nm)
325+
assert_array_equal(self.f & self.im, False)
313326
assert_array_equal(self.nm & True, self.nm)
314327
assert_array_equal(self.im & False, self.f)
328+
assert_array_equal(True & self.nm, self.nm)
329+
assert_array_equal(False & self.im, self.f)
315330
assert_array_equal(self.nm | self.t, self.t)
316331
assert_array_equal(self.im | self.f, self.im)
332+
assert_array_equal(self.t | self.nm, self.t)
333+
assert_array_equal(self.f | self.im, self.im)
317334
assert_array_equal(self.nm | True, self.t)
318335
assert_array_equal(self.im | False, self.im)
336+
assert_array_equal(True | self.nm , self.t)
337+
assert_array_equal(False | self.im, self.im)
319338
assert_array_equal(self.nm ^ self.t, self.im)
320339
assert_array_equal(self.im ^ self.f, self.im)
340+
assert_array_equal(self.t ^ self.nm, self.im)
341+
assert_array_equal(self.f ^ self.im, self.im)
321342
assert_array_equal(self.nm ^ True, self.im)
322343
assert_array_equal(self.im ^ False, self.im)
344+
assert_array_equal(True ^ self.nm, self.im)
345+
assert_array_equal(False ^ self.im, self.im)
346+
347+
# test scalar array specialization
348+
norm = np.zeros(114, dtype=bool)[1:]
349+
scalar_f = np.ndarray(buffer=np.zeros((), dtype=bool), strides=(0,),
350+
shape=norm.shape, dtype=bool)
351+
scalar_t = np.ndarray(buffer=np.ones((), dtype=bool), strides=(0,),
352+
shape=norm.shape, dtype=bool)
353+
assert_array_equal(scalar_f & scalar_t, norm)
354+
assert_array_equal(scalar_t & scalar_t, ~norm)
355+
assert_array_equal(scalar_f & scalar_f, norm)
356+
assert_array_equal(scalar_f | scalar_t, ~norm)
357+
assert_array_equal(scalar_t | scalar_t, ~norm)
358+
assert_array_equal(scalar_f | scalar_f, norm)
359+
assert_array_equal(scalar_f ^ scalar_t, ~norm)
360+
assert_array_equal(scalar_t ^ scalar_t, norm)
361+
assert_array_equal(scalar_f ^ scalar_f, norm)
323362

324363

325364
class TestBoolCmp(TestCase):

0 commit comments

Comments
 (0)
0