Proper way to accumulate results

[gcesars]

I have a 3 dimension sparse field, and I am able to run several processes with them at an impressive speed, much faster than tf.function.
I was experimenting with some code before moving to the next steps, and the simple operation of counting the number of active cells in the sparse tree is weirdly slow; it takes 1000x more than the much more complex operations I have done before.

Example dummy code below.

I assume the issue are the += operations making the thread to stop and wait, but I can't give any proof.

Anyway, what is the recommended way to design accumulator loops with good performance?

values[0] = 0
values[1] = 0
values[2] = 0
for i in ti.ndrange(self.dim_1.shape[0]):
    if ti.is_active(self.dim_1, i):
        values[0] += 1
        for j in ti.ndrange(self.dim_2.shape[1]):
            if ti.is_active(self.dim_2, [i, j]):
                values[1] += 1
            for k in ti.ndrange(self.dim_3.shape[2]):
                if ti.is_active(self.dim_3, [i, j, k]):
                    values[2] += 1