[MPS] On MacOS-26 torch.full fails for 4+Gb tensors #161265
Description
🐛 Describe the bug
If one runs following example on macOS Tahoe
import torch
a = torch.ones(2, (1 << 31) + 5, dtype=torch.int8, device='mps')
print(a[1, -2])
print(a[:, -2])result will be
tensor(0, device='mps:0', dtype=torch.int8)
tensor([1, 1], device='mps:0', dtype=torch.int8)
But should have been 1 in both cases
Reason for that are relying on unsupported use of fillBuffer that for some reason stopped working for length exceeding 4Gb
pytorch/aten/src/ATen/mps/MPSStream.mm
Line 161 in 639b8cc
I.e. following Swift script returns non-unfirom values on MacOS Tahoe
import Metal
let bufferSize = 1<<32 + 4
guard let device = MTLCreateSystemDefaultDevice() else { fatalError("No Metal device found") }
guard let buffer = device.makeBuffer(length: bufferSize, options: .storageModeShared) else { fatalError("Failed to create buffer") }
guard let cmdQueue = device.makeCommandQueue() else { fatalError("Failed to create command queue") }
guard let cmdBuffer = cmdQueue.makeCommandBuffer() else { fatalError("Failed to create command buffer") }
guard let blitEncoder = cmdBuffer.makeBlitCommandEncoder() else { fatalError("Failed to create blit encoder") }
blitEncoder.fill(buffer: buffer, range: 0..<bufferSize, value: 0x42)
blitEncoder.endEncoding()
cmdBuffer.commit()
cmdBuffer.waitUntilCompleted()
let tailOffs = 8
let hostPtr = buffer.contents().bindMemory(to: UInt8.self, capacity: bufferSize)
let tail = Array(UnsafeBufferPointer(start: hostPtr + (bufferSize - tailOffs), count: tailOffs))
for (idx, val) in tail.enumerated() {
print("Offs 0x\(String(bufferSize - tailOffs + idx, radix: 16)): 0x\(String(val, radix: 16))")
}Versions
2.8.0, nightly