Inefficient count_nonzero
Hi @ngc92 ,

I have replicated the reported behaviour on CPU runtime as per attached [cpu-gist](https://colab.research.google.com/gist/SuryanarayanaY/39d9263b1b84cd59ddad059140e29ae6/60772_cpu.ipynb).

But with GPU runtime the case is totally different. With `uint32` as dtype actually it is taking much more time than `int64` since some Ops with `unit32` as dtype are not supported on GPU . It seems this API implementation supports  `int64` for best performance. I tried with `int32` also where performance is not good at all compared to `int64` even though all Ops ran on GPU only. I have attached all the results in the attached [gpu-gist](https://colab.research.google.com/gist/SuryanarayanaY/2d9a40aeb2c35f2ae8363b8d58c8a16e/60772_gpu.ipynb).

Please go through the results and shave your views. Thanks!
 @SuryanarayanaY 
Thanks for the replication. 

First, your results seem to confirm that skipping the redundant comparison with zero leads to a measureable improvement in performance in both cases, so that should be a relatively straightforward improvement.

The missing GPU kernels make the other improvements more tricky, though.

I've now run the current `count_nonzeros` implementation also on GPU under the tf profiler, with the following results:
```
Kernel Name                           Op Name                 Total    Avg     Min     Max
Cast_GPU_DT_BOOL_DT_INT64_kernel      count_nonzero/Cast      12,488   4,162   4,018   4,303
cub::DeviceSegmentedReduceKernel      count_nonzero/Sum       9,973    3,324   3,249   3,470
NotEqual_GPU_DT_BOOL_DT_BOOL_kernel   count_nonzero/NotEqual  7,768    2,589   2,429   2,907
```
As you can see, the conversion to int64 actually takes *more* time than the reduction itself.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Performance

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

2.12

### Custom Code

Yes

### OS Platform and Distribution

_No response_

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

The current implementation of `tf.math.count_nonzero` is extremely inefficient, when trying to count how often a boolean condition is true in a large tensor. This is because the operation first converts the boolean tensor into `tf.int64`, before feeding it to `tf.reduce_sum`. As that is an operation that is entirely bandwidth-bound, instead of transferring `n` bytes for an input of `n` elements, now it has to transfer `n + 8n` for the cast operation, and an additional `8n` for the reduction, making this 17 times (!!) more inefficient than it needs to be.

The benchmark below show the following:
1) Plain `count_nonzero`
2) The same implementation as count_nonzero, except that a superfluous comparison with zero is omitted (which is needed for other data types to convert to bool, but appears to not be optimized away even if the tensor is already of bool type). This gives a small speed-up
3) Use `uint32` as accumulation type, instead of `int64`. This shows the tremendous effect of bandwidth.


Here are a few suggestions for improvements, in increasing order of difficulty:
1) If the input is already of bool type, there is no need to do `gen_math_ops.not_equal(input, zero)`
2) if the input is of bool type and has less than 2**32 elements, use `uint32` as the accumulator, instead of `int64`
3) implement a dedicated `count_nonzero` op that directly performs the reduction on boolean tensors
4) implement a generic `reduce_sum` that can have a different accumulator than the input type. This could also be helpful, e.g., for summing up lots of float/half values, where one would like to keep the error in check by using double/float accumulation.

If there is interest, I can provide a patch for 1 and 2. I'm not sure about 3, as I guess this new op would need to be implemented for CPU, GPU and TPU for this to make sense.

### Standalone code to reproduce the issue

```shell
Here is a very simple benchmark that illustrates the issue:

import tensorflow as tf
import time

print("TF VERSION", tf.version.VERSION)

data = tf.random.uniform(shape=(32, 1000000))
predicate = tf.greater(data, 0.5)

start = time.time()

for _ in range(100):
    a = tf.math.count_nonzero(predicate, axis=1)

print(time.time() - start)


start = time.time()

for _ in range(100):
    a = tf.math.reduce_sum(tf.cast(predicate, tf.int64), axis=1)

print(time.time() - start)


start = time.time()

for _ in range(100):
    a = tf.math.reduce_sum(tf.cast(predicate, tf.uint32), axis=1)

print(time.time() - start)
```


### Relevant log output

```shell
TF VERSION 2.12.0
4.673052549362183
4.253457069396973
2.3325610160827637
```
</details>