Optimizing cpu kernels for reductions #484
Conversation
| /// | ||
| /// This makes all the reduced elements sequential when indexed using [NdIndex]. | ||
| #[inline(always)] | ||
| pub(crate) fn index_for_reductions<S: Shape, Ax: Axes>( |
There was a problem hiding this comment.
@nkoppel i'm pretty sure this does exactly what the cuda version below does. Thoughts on if we could combine them? I'd like to keep the cpu one not using vecs, which I think is the biggest difference.
There was a problem hiding this comment.
I think we can factor out the shared logic from index_for_reductions and permute_for_reductions, but I'd prefer to keep the filtering of broadcasted dimensions in permute_for_reductions, to reduce overhead from get_strided_index for broadcasted tensors.
This implementation looks correct to me, and is likely faster than my implementation, so go ahead and use yours.
| new_shape[num_non_reduced_dims + i_reduced] = dims[i_src]; | ||
| new_strides[num_non_reduced_dims + i_reduced] = strides[i_src]; | ||
| i_reduced += 1; |
There was a problem hiding this comment.
An optimization for this is probably to sort the strides in this section of new_shape/new_strides by decreasing size, to maximize locality when indexing.
This moves cpu kernels closer to what cuda kernels do for reductions:
Additionally this adds some shortcuts for full reductions (to 0d)