Offsets on mps kernel indexing

arunppsg · August 8, 2023, 6:58pm

I am trying to understand kernel indexing for MPS backend in pytorch. I have two questions:

Though kernel_index_offsets recieves num_offsets argument, it is not using it. Am I overlooking something here?

pytorch/pytorch/blob/2764ead4296541b324ae0e8e2b6200eec0e90c4e/aten/src/ATen/mps/IndexKernels.h#L211-L225


      
          kernel void kernel_index_offsets(constant packed_uint3 * strides         [[buffer(0)]],
                                           device uint3          * data_offsets    [[buffer(1)]],
                                           constant uint         * iter_shape      [[buffer(2)]],
                                           constant uint         & num_dimensions  [[buffer(3)]],
                                           constant uint         & num_offsets     [[buffer(4)]],
                                           uint thread_index [[thread_position_in_grid]]) {
              data_offsets[thread_index] = 0;
              uint32_t idx = thread_index;
              for (uint32_t dim = 0; dim < num_dimensions; dim++) {
                  uint32_t remainder = idx % iter_shape[dim];
                  idx /= iter_shape[dim];
          
                  data_offsets[thread_index] += remainder * strides[dim];
              }
          }

Why is nOffsets assigned 3 in the following line:

github.com

pytorch/pytorch/blob/2764ead4296541b324ae0e8e2b6200eec0e90c4e/aten/src/ATen/native/mps/operations/BinaryKernel.mm#L179-L179


      
          constexpr uint32_t nOffsets = 3;

Apologies if this question is out of place here. I am interested in contributing to pytorch (in this General MPS op coverage tracking issue · Issue #77764 · pytorch/pytorch · GitHub) and working towards getting an understanding of the codebase.

qqaatw · August 13, 2023, 7:28am

Hi,

For your questions:

Yes it is indeed not being used as this kernel assumes the number of offsets being 3 (you can see the data type of data_offsets, which is uint3).
The kernel was originally designed for handling binary operations, i.e. binary inputs, so the offsets represent input(self), other, and output, respectively.