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.

Topic		Replies	Views
MPS working group? hardware-backends	2	368	March 12, 2024
About "[FSDP] Don't move ignored params / buffers to device"	1	222	August 2, 2024
Memory operations on a custom backend hardware-backends	4	1147	July 5, 2022
How to approach targets that don't support i64/f64?	1	795	November 20, 2022
On input arguments to dispatch stub	0	405	August 6, 2023

Offsets on mps kernel indexing

Related topics