Generalize softmax for packed dim vs non packed dim (#5755)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/5755

## Context

This diff performs a rewrite of the `softmax` shaders. Previously, the shaders were separated into the `channels` case and `batch_height_width` case. This is because channels packing was the only packing format used when these shaders were written, thus these shaders represented the case when the reduction dim is equal to the packed dim, and the case when the reduction dim is orthogonal to the packed dim respectively.

Now that we expect tensors to be width packed as well, the `channels`/`batch_height_width` separation no longer makes sense. This diff consolidates both cases into a single shader that takes the `packed_dim` and `reduce_dim` as specialization constants, and selects the correct function to execute based on if they are the same or different.

Additionally, I implemented a optimization in the form of using a co-operative algorithm to allow multiple threads to co-operate in computing max and sum.

More details can be found in the comments of the new shader file.
ghstack-source-id: 245571371

Reviewed By: jorgep31415

Differential Revision: D63642091

fbshipit-source-id: cfe960a20cacdb7670390f7626fbf64366a734b0