Slow operations using bool directly than casting to int first for batch of small arrays?

[JoeyTeng]

I am trying to find the fastest way of performing logical and on 3 operands over a huge batch of data. I benchmarked several approaches and find that on GPUs (T4), using bool arrays directly are significantly slower (15%+) than converting them to int arrays then do reduction. However, this behaviour is not observed on more advanced GPUs (V100), and is reversed on TPU (ints are slower). Why is this the case?

The benchmark code are as below. (Colab available here)

Benchmark Code

import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.random as random

keyA, keyB, keyC = random.split(random.PRNGKey(20230609), 3)
boolsA = random.bernoulli(random.PRNGKey(20230609), p=0.011, shape=(64 * 960, 1000)).block_until_ready()
boolsB = random.bernoulli(random.PRNGKey(20230609), p=0.312, shape=(64 * 960, 1000)).block_until_ready()
boolsC = random.bernoulli(random.PRNGKey(20230609), p=0.999, shape=(64 * 960, 1000)).block_until_ready()

_logical_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.logical_and(a, jnp.logical_and(b, c))))).lower(boolsA, boolsB, boolsC).compile()
_bitwise_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: lax.bitwise_and(a, lax.bitwise_and(b, c))))).lower(boolsA, boolsB, boolsC).compile()
_operand_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: a & b & c))).lower(boolsA, boolsB, boolsC).compile()
_all_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).all()))).lower(boolsA, boolsB, boolsC).compile()

_prod_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).astype(jnp.uint8).prod() != 0))).lower(boolsA, boolsB, boolsC).compile()
_prod_weak_type_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).astype(int).prod() != 0))).lower(boolsA, boolsB, boolsC).compile()

_min_bool_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).min()))).lower(boolsA, boolsB, boolsC).compile()
_min_int_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).astype(int).min() > 0))).lower(boolsA, boolsB, boolsC).compile()
_min_int2_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).astype(int).min() != 0))).lower(boolsA, boolsB, boolsC).compile()
_min_int3_ = jax.jit(jax.vmap(jax.vmap(lambda a, b, c: jnp.array((a, b, c)).astype(jnp.uint8).min() != 0))).lower(boolsA, boolsB, boolsC).compile()

print("JAX Version:", jax.__version__)
print("Python Version:", python_version())
!cat /var/colab/hostname

print("Bool operations")
%timeit _logical_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _bitwise_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _operand_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _all_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _min_bool_(boolsA, boolsB, boolsC).block_until_ready()

print("min among casted int array")
%timeit _min_int_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _min_int2_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _min_int3_(boolsA, boolsB, boolsC).block_until_ready()

print("prod reduction along a casted int array")
%timeit _prod_(boolsA, boolsB, boolsC).block_until_ready()
%timeit _prod_weak_type_(boolsA, boolsB, boolsC).block_until_ready()

The outputs are as below

T4, JAX 0.4.10, Python 3.10.12

JAX Version: 0.4.10
Python Version: 3.10.12
gpu-t4-s-m0yjo20dx7m8
Bool operations
1.4 ms ± 112 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.26 ms ± 16.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.28 ms ± 5.61 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.28 ms ± 7.95 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.29 ms ± 3.18 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
min among casted int array
1.14 ms ± 21.9 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.13 ms ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.11 ms ± 5.29 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
prod reduction along a casted int array
1.12 ms ± 5.52 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.12 ms ± 9.48 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

V100, JAX 0.4.10, Python 3.10.12

JAX Version: 0.4.10
Python Version: 3.10.12
gpu-v100-s-21gphkyu3p4z6
Bool operations
481 µs ± 16.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
486 µs ± 6.04 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
482 µs ± 5.58 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
491 µs ± 9.99 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
498 µs ± 13.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
min among casted int array
493 µs ± 4.87 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
507 µs ± 7.67 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
497 µs ± 6.64 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
prod reduction along a casted int array
500 µs ± 7.75 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
500 µs ± 10.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

TPU, JAX 0.3.25, Python 3.10.12

JAX Version: 0.3.25
Python Version: 3.10.12
tpu-s-3g1uc5qc0i4bk
Bool operations
3.42 ms ± 174 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.88 ms ± 73.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.95 ms ± 95.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.76 ms ± 138 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
2.89 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
min among casted int array
5.48 ms ± 405 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.96 ms ± 281 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.65 ms ± 178 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
prod reduction along a casted int array
8.55 ms ± 782 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.37 ms ± 187 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)