Leveraging Unified Memory for MPS Tensors on Apple Silicone

### 🚀 The feature, motivation and pitch

### Background

In the current PyTorch MPS backend, tensor backing buffers are allocated in the owner device’s private memory. When a tensor is moved between cpu and mps, PyTorch:
1. Allocates a new backing buffer in the destination device’s private memory.
2. Encodes a Metal `blit` operation to copy data from source to destination.

This behavior is consistent and correct, but may not fully leverage the unified physical memory architecture of Apple Silicon. Specifically in cases where the source and destination tensors are not modified, this results in two sets of identical buffers representing the same tensor on source and destination device. Our proposal to address this issue is comprised of three parts:
1. Allocate all tensors in unified memory instead of device's private memory.
2. In case of moving read-only tensors: Since the data cannot be mutated, avoid duplication and `memcpy` of the source backing buffer, and instead reference the same unified memory for the destination tensor.
3. In case of assigning the result of a tensor move to itself: Since source tensor and its backing buffer are discarded after device move, avoid duplication and `memcpy` of the source backing buffer, and instead reference the same unified memory for the destination tensor.

### Proposal Part 1 - Allocate Tensors in Unified Memory
Part 1 of the proposal is implemented and currently going through testing and verification. Should Parts 2 & 3 be rejected, Part 1 can still be merged, as it is independent of the subsequent parts and can still provide performance benefits.

**Summary**
Instead of allocating tensor backing buffers in device-private memory by default, we propose allocating _all_ tensors in the unified (shared) memory accessible by both CPU and MPS. Under this model:

- Tensor device transfers (cpu ↔ mps) would still create a destination tensor with new ownership.
- However, data movement would use a `memcpy` on unified memory plus a memory barrier, rather than a Metal cpu ↔ mps `blit` operation where applicable.
- This part of the proposal on its own is **not** intended to eliminate destination allocations for `.to(device=...);` calls rather, it focuses on reducing the cost of data movement, when both source and destination backing buffers reside in unified memory. It also paves the way for the next parts of the proposal to specifically reduce memory duplication.

**Expected Impact**
- Memory allocation size: No significant reduction expected in the general case.
- Transfer cost: Potential reduction by replacing CPU ↔ GPU `blit` operations with lighter-weight `memcpy` calls on unified memory.
- Semantics: No change to PyTorch tensor ownership or device semantics.

**Code Changes**
- Tensor Allocation [[EmptyTensor.cpp](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/mps/EmptyTensor.cpp#L21)]
Tensor allocation is performed by fetching an allocator class through the `GetMPSAllocator` call. If device supports unified memory, we propose to pass `useSharedAllocator = true` for allocator creation. As a result, all tensors allocated through this allocator will reside in unified memory.

- Tensor Device Move [[Copy.mm](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/mps/operations/Copy.mm#L161)]
Currently, during an CPU → MPS copy, an explicit blit (`stream→copy_and_sync()`) is encoded to move bytes from CPU-backed private memory into MPS. This will no longer be needed as both source and destination will reside in shared memory. Instead, we propose to perform a `memcpy` of the backing buffer on cpu-visible unified memory. Since a mere `memcpy` does not carry tensor-specific information, we also manually perform cpu-side negation/conjugation if necessary, as well as to manually change ownership to `device(at::kMPS)`. We also add a memory barrier (`stream->synchronize(SyncType::COMMIT_AND_WAIT);`) to ensure writes to the tensor by source device are visible to the destination device. Please note that at this stage memory duplication still occurs but it is on unified memory rather than device private memory.

### Proposal Part 2 - Avoid Duplication for Read-only Tensors
**Summary**
With Part 1 implemented, all tensors will reside in unified memory and be accessible from both CPU and MPS. However, when a tensor is moved between devices, we cannot enforce source and destination tensors to reference the same backing buffer, as PyTorch semantics require them to behave independently, i.e. write to either tensor should not modify the other. This problem does not however exist in case of read-only tensors, and we can reference the same memory with new device ownership instead of maintaining two copies of the same data.

**Code Changes**
Still in progress and subject to approval of the proposal.

**Expected Impact**
- Memory allocation size: Reduction by 50% for read-only tensors, by maintaining a single copy of the backing buffer in unified memory for both source and destination tensors.
- Transfer cost: Reduction by avoiding a `memcpy` on unified memory all together, and simply changing the tensor ownership.
- Semantics: Mutable tensors will still maintain two copies for source and destination tensors. Therefore, no change to PyTorch tensor ownership or device semantics.

### Proposal Part 3 - Avoid Duplication when Assigning a Tensor Device Move to Itself
**Summary**
Currently and with Part 1 of the proposal in place, when the result of a tensor device move is assigned to itself (`a = a.to(device='mps')`), a duplicate destination buffer is allocated and copied from source, and immediately after the source tensor is discarded and its backing buffer freed. This is equivalent to simply moving the backing buffer from one place of the unified memory to another. Thus we propose to eliminate the unnecessary duplication, and simply reference the same source backing buffer with new ownership for the destination tensor.

**Code Changes**
Still in progress and subject to approval of the proposal. We are also still investigating if it is possible to distinguish when a tensor device move is assigned to a new tensor vs. when it is assigned to itself.

**Expected Impact**
- Memory allocation size: overall memory consumption should remain unchanged, but peak memory consumption has the potential to be reduced by reusing the source backing buffer.
- Transfer cost: Reduction by avoiding a `memcpy` on unified memory all together, and simply changing the tensor ownership.
- Semantics: If both source and destination tensors remain alive, they will still maintain two copies for source and destination tensors. Therefore, no change to PyTorch tensor ownership or device semantics.

### Alternatives

Instead of Parts 2 & 3 of the proposal, a general Copy-On-Write implementation was also considered that would keep both source and destination tensors referencing the same backing buffer in unified memory, until either one was modified. However, as we understand, [the last Copy-On-Write attempt](https://docs.google.com/document/d/10fyjbbpTc9-7T8QHIkPtJxnHyqMzqD4DhG6vppz8ujc/edit?tab=t.0) was ultimately rejected due to the design being too brittle. Thus, we propose to reduce the optimization scope to read-only tensors and self-assigning device moves.

### Additional context

- Part 1 of the proposal is implemented and currently at a 97% pass rate for `tests_mps.py`. We are tracking down the remaining 3% failures.
- Parts 2 & 3 are still under investigation and not implemented. We hope to reach an understanding on whether or not these parts of the proposal in theory are acceptable to the maintainers before attempting to implement.

cc @kulinseth @malfet @DenisVieriu97 @jhavukainen