PR #181949: [MPS] Enhance Col2Im tensor op to avoid .contiguous() call and make corresponding Metal kernel stride-aware

• Repository: pytorch/pytorch
• Author: bornaehsani
• State: open | merged: False
• Link: https://github.com/pytorch/pytorch/pull/181949

Description (problem / solution / changelog)

Col2Im Tensor OP Enhancement

• Removing the .contiguous() call on col_tensor in col2im_out_mps_template() operation
• Updating col2im_kernel() kernel to take in inner strides and manually calculate strided col_index

Benchmarked using following script

import torch
import torch.nn.functional as F
import torch.utils.benchmark as benchmark

# col2im_out_mps_template() calls input.contiguous() before PR #181949.
# Transpose dims 1 and 2 to make the input non-contiguous and force the copy.
x = torch.randn(32, 16384, 27, device="mps").transpose(1, 2)
assert not x.is_contiguous()

m = benchmark.Timer(
    stmt="F.fold(x, output_size=(128, 128), kernel_size=3, padding=1, stride=1)",
    globals={"F": F, "x": x},
    num_threads=1,
).blocked_autorange(min_run_time=10.0)

print(f"col2im: mean={m.mean * 1e6:.2f} us  median={m.median * 1e6:.2f} us  iqr={m.iqr * 1e6:.2f} us")

On an M3 Macbook Pro with 48GB memory we observe a 3.1x speedup.

Issue #181946

Changed files

• aten/src/ATen/native/mps/kernels/Col2Im.metal (modified, +7/-4)
• aten/src/ATen/native/mps/operations/Col2Im.mm (modified, +6/-2)

PR #181951: [MPS] Enhance HistogramKernel tensor op to avoid unnecessary .contiguous() call

• Repository: pytorch/pytorch
• Author: bornaehsani
• State: open | merged: False
• Link: https://github.com/pytorch/pytorch/pull/181951

Description (problem / solution / changelog)

• Removing the .contiguous() call for bin_edges and the is_contiguous() assertion in histogramdd_out_mps_template()
• The tensor is never passed into kernel, it is only consumed on CPU through the .item() which already supports non-contiguous tensors.

Benchmarked using following script:

import torch
import torch.nn.functional as F
import torch.utils.benchmark as benchmark

bins_buf = torch.linspace(-4.0, 4.0, 514, device=self._device)
bins = bins_buf[::2]  # 257 edges → 256 bins, stride=2, non-contiguous
self.assertFalse(bins.is_contiguous())
x = torch.randn(1_000_000, device=self._device)

m = benchmark.Timer(
    stmt="torch.histogram(x, bins=bins)",
    globals={"torch": torch, "x": x, "bins": bins},
    num_threads=1,
).blocked_autorange(min_run_time=10.0)

print(f"histogramkernel: mean={m.mean * 1e6:.2f} us  median={m.median * 1e6:.2f} us  iqr={m.iqr * 1e6:.2f} us")

On an M3 Macbook Pro with 48GB memory we observe a 1.08x speedup.

Issue #181946

Changed files

• aten/src/ATen/native/mps/operations/HistogramKernel.mm (modified, +1/-7)