[MPS] Use metal shaders for all view ops (pytorch#143375)

Before this PR Metal  shaders were used to scatter/gather 1-5 dimensional tensors.
This PR introduces generalized ones that could be used for any dimensionality and as results  gets rid of 700+ lines complex and untested code that might not even work as expected.
Generalized gather shader looks as follows
```metal
kernel void gather_kernel_n(uint linear_index           [[thread_position_in_grid]],
                            constant void * src_        [[buffer(0)]],
                            device void * dst_          [[buffer(1)]],
                            constant uint32_t * size    [[buffer(2)]],
                            constant uint32_t * stride  [[buffer(3)]],
                            constant uint32_t & numel   [[buffer(4)]],
                            constant int32_t & ndim     [[buffer(5)]]) {{
    if (linear_index >= numel) return;

    constant {0} * src = (constant {0} *)src_;
    device {1} * dst = (device {1} *)dst_;

    uint64_t src_offs = 0;
    auto src_idx = linear_index;
    for(int dim = ndim - 1; dim >= 0; --dim) {{
      src_offs += stride[dim] * (src_idx % size[dim]);
      src_idx /= size[dim];
    }}

    dst[linear_index] = cast<{1}>(src[src_offs]);
}}
```

Which, according to the following benchmark
```python
from timeit import default_timer

import torch
import torch.utils.cpp_extension
from torch.utils.benchmark import Measurement, Timer

t = Timer(
    stmt=f"y.copy_(x);torch.mps.synchronize()",
    setup=f"x=torch.rand(4, 5, 16, 64, 33, 24, dtype=torch.float32, device='mps')[:,:,:,:24,:24,];y=torch.empty(x.shape, device=x.device, dtype=x.dtype)",
    language="python", timer=default_timer
)
print(t.blocked_autorange())
```
Is almost twice as fast as previous implementation (i.e. on Mac Book M2 Pro it returns 2.9ms for MPS version vs 1.5ms for shader one

On MacOS Sequoia [`gatherWithUpdatesTensor: indicesTensor:...`](https://developer.apple.com/documentation/metalperformanceshadersgraph/mpsgraph/gather(withupdatestensor:indicestensor:axis:batchdimensions:name:)?language=objc) crashes if invoked with complex data type, as one can see by running the code below
```swift
import Metal
import MetalPerformanceShadersGraph

func gatherComplexMPS(device: MTLDevice,
                inp_buf: MTLBuffer, idx_buf: MTLBuffer,
                out_buf: MTLBuffer,
                inp_elem: Int, upd_elem: Int) {
  let graph = MPSGraph()
  let inputPlaceholder = graph.placeholder(shape: [inp_elem as NSNumber], dataType: .complexFloat32, name: nil)
  let indicesPlaceholder = graph.placeholder(shape: [upd_elem as NSNumber], dataType: .int64, name: nil)
  let outNode = graph.gather(withUpdatesTensor: inputPlaceholder, indicesTensor: indicesPlaceholder, axis: 0, batchDimensions: 0, name: nil)
  let mpsInputBuffer = MPSGraphTensorData(inp_buf, shape: [inp_elem as NSNumber], dataType: .complexFloat32)
  let mpsIndicesBuffer = MPSGraphTensorData(idx_buf, shape: [upd_elem as NSNumber], dataType: .int64)
  let mpsOutputBuffer = MPSGraphTensorData(out_buf, shape: [inp_elem as NSNumber], dataType: .complexFloat32)
  guard let queue = device.makeCommandQueue() else { fatalError("Can't make queue") }
  graph.run(with: queue, feeds: [inputPlaceholder: mpsInputBuffer,
                               indicesPlaceholder: mpsIndicesBuffer ],
            targetOperations: nil, resultsDictionary: [outNode: mpsOutputBuffer])
}

func makeBufferWithValues<T>(device: MTLDevice, values: [T]) -> MTLBuffer {
  guard let buf = device.makeBuffer(length: values.count * MemoryLayout<T>.size, options: [.storageModeShared]) else { fatalError("Can't alloc") }
  let buf_data = buf.contents().assumingMemoryBound(to: T.self)
  for i in 0..<values.count {
    buf_data[i] = values[i]
  }
  return buf
}

guard let device = MTLCopyAllDevices().first else { fatalError("Not Metal device found") }
print("Using device \(device.name)")

let inp_buf = makeBufferWithValues(device: device, values: [1.0, 2.0 , 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
let idx_buf = makeBufferWithValues(device: device, values: [0, 1, 2, 3])
guard let out_buf = device.makeBuffer(length:8 * MemoryLayout<Float>.size, options: [.storageModeShared]) else { fatalError("Can't alloc") }

gatherComplexMPS(device: device, inp_buf: inp_buf, idx_buf: idx_buf, out_buf: out_buf, inp_elem: 4, upd_elem: 4)
```

Fixes pytorch#143140
Pull Request resolved: pytorch#143375
Approved by: https://github.com/albanD

Author: malfet

aten/src/ATen/mps/IndexKernels.h

@@ -3,10 +3,6 @@
 namespace at::mps {
 
 static const char* SCATTER_OPS_TEMPLATE = R"METAL_SCATTER(
-struct __attribute__ ((packed)) packed_uint5{{
-  uint32_t x; uint32_t y; uint32_t z; uint32_t w; uint32_t u;
-}};
-
 template<typename Y, typename X>
 Y cast(const X x);
 
@@ -15,32 +11,26 @@ template<>
  return {2};
 }}
 
-kernel void scatter_kernel_5(uint linear_index              [[thread_position_in_grid]],
-                             constant void * src_           [[buffer(0)]],
-                             device void * dst_             [[buffer(1)]],
-                             constant packed_uint5 & size   [[buffer(2)]],
-                             constant packed_uint5 & stride [[buffer(3)]],
-                             constant uint32_t & numel      [[buffer(4)]]) {{
+kernel void scatter_kernel_n(uint linear_index          [[thread_position_in_grid]],
+                             constant void * src_       [[buffer(0)]],
+                             device void * dst_         [[buffer(1)]],
+                             constant uint32_t * size   [[buffer(2)]],
+                             constant uint32_t * stride [[buffer(3)]],
+                            constant uint32_t & numel   [[buffer(4)]],
+                            constant int32_t & ndim     [[buffer(5)]]) {{
     if (linear_index >= numel) return;
 
     constant {0} * src = (constant {0} *)src_;
     device {1} * dst = (device {1} *)dst_;
 
-    packed_uint5 local_index;
-    local_index.x = linear_index / (size.u * size.w * size.z * size.y) % size.x;
-    local_index.y = linear_index / (size.u * size.w * size.z) % size.y;
-    local_index.z = linear_index / (size.u * size.w) % size.z;
-    local_index.w = linear_index / size.u % size.w;
-    local_index.u = linear_index % size.u;
-
-    packed_uint5 strided_index;
-    strided_index.x = local_index.x * stride.x;
-    strided_index.y = local_index.y * stride.y;
-    strided_index.z = local_index.z * stride.z;
-    strided_index.w = local_index.w * stride.w;
-    strided_index.u = local_index.u * stride.u;
-
-    dst[strided_index.x + strided_index.y + strided_index.z + strided_index.w + strided_index.u] = cast<{1}>(src[linear_index]);
+    uint64_t dst_offs = 0;
+    auto dst_idx = linear_index;
+    for(int dim = ndim - 1; dim >= 0; --dim) {{
+      dst_offs += stride[dim] * (dst_idx % size[dim]);
+      dst_idx /= size[dim];
+    }}
+
+    dst[dst_offs] = cast<{1}>(src[linear_index]);
 }}
 
 kernel void scatter_kernel_4(uint linear_index              [[thread_position_in_grid]],
@@ -121,10 +111,6 @@ kernel void scatter_kernel_1(uint linear_index              [[thread_position_in
 )METAL_SCATTER";
 
 static const char* GATHER_OPS_TEMPLATE = R"METAL_GATHER(
-struct __attribute__ ((packed)) packed_uint5{{
-  uint32_t x; uint32_t y; uint32_t z; uint32_t w; uint32_t u;
-}};
-
 template<typename Y, typename X>
 Y cast(const X x);
 
@@ -133,33 +119,26 @@ template<>
  return {2};
 }}
 
-kernel void gather_kernel_5(uint linear_index               [[thread_position_in_grid]],
-                            constant void * src_            [[buffer(0)]],
-                            device void * dst_              [[buffer(1)]],
-                            constant packed_uint5 & size    [[buffer(2)]],
-                            constant packed_uint5 & stride  [[buffer(3)]],
-                            constant uint32_t & numel       [[buffer(4)]]) {{
+kernel void gather_kernel_n(uint linear_index           [[thread_position_in_grid]],
+                            constant void * src_        [[buffer(0)]],
+                            device void * dst_          [[buffer(1)]],
+                            constant uint32_t * size    [[buffer(2)]],
+                            constant uint32_t * stride  [[buffer(3)]],
+                            constant uint32_t & numel   [[buffer(4)]],
+                            constant int32_t & ndim     [[buffer(5)]]) {{
     if (linear_index >= numel) return;
 
     constant {0} * src = (constant {0} *)src_;
     device {1} * dst = (device {1} *)dst_;
 
+    uint64_t src_offs = 0;
+    auto src_idx = linear_index;
+    for(int dim = ndim - 1; dim >= 0; --dim) {{
+      src_offs += stride[dim] * (src_idx % size[dim]);
+      src_idx /= size[dim];
+    }}
 
-    packed_uint5 local_index;
-    local_index.x = linear_index / (size.u * size.w * size.z * size.y) % size.x;
-    local_index.y = linear_index / (size.u * size.w * size.z) % size.y;
-    local_index.z = linear_index / (size.u * size.w) % size.z;
-    local_index.w = linear_index / size.u % size.w;
-    local_index.u = linear_index % size.u;
-
-    packed_uint5 strided_index;
-    strided_index.x = local_index.x * stride.x;
-    strided_index.y = local_index.y * stride.y;
-    strided_index.z = local_index.z * stride.z;
-    strided_index.w = local_index.w * stride.w;
-    strided_index.u = local_index.u * stride.u;
-
-    dst[linear_index] = cast<{1}>(src[strided_index.x + strided_index.y + strided_index.z + strided_index.w + strided_index.u]);
+    dst[linear_index] = cast<{1}>(src[src_offs]);
 }}
 
 kernel void gather_kernel_4(uint linear_index               [[thread_position_in_grid]],

aten/src/ATen/native/mps/OperationUtils.h

@@ -81,10 +81,6 @@ std::string getArrayRefString(const IntArrayRef s);
 // use has_storage() on the returned tensor to determine if src actually is a view
 Tensor gatherViewTensor(const Tensor& src, Tensor& dst);
 Tensor& scatterViewTensor(const Tensor& src, Tensor& output);
-bool canSliceViewTensor(const TensorBase& src, MPSShape* mpsShape);
-MPSGraphTensorData* getMPSGraphTensorDataForView(const TensorBase& src,
-                                                 MPSShape* mpsShape,
-                                                 const MPSDataType mpsDataType);
 MPSGraphTensor* castToIHFTypes(MPSGraph* mpsGraph,
                                MPSGraphTensor* inputTensor,
                                const TensorBase& input,