Pytorch CPU vs GPU (Apple Silicon MPS) Performance Comparison (Scaling with Matrix Size)

pytorch_cpu_mps_comparison.py

import torch
import timeit

def compare_cpu_mps(size, dimensions = 2, iterations = 2_500):
    a_cpu = torch.rand(size, device='cpu')
    b_cpu = torch.rand(tuple((size for n in range(dimensions))), device='cpu')
    a_mps = torch.rand(size, device='mps')
    b_mps = torch.rand(tuple((size for n in range(dimensions))), device='mps')

    print('    - cpu', timeit.timeit(lambda: a_cpu @ b_cpu, number=iterations))
    print('    - mps', timeit.timeit(lambda: a_mps @ b_mps, number=iterations))

for n in [100, 250, 500, 1000]:
    print(f"With {n}:")
    for d in [1, 2, 3]:
        print(f"  In {d}D:")
        compare_cpu_mps(n, d)

        

result.txt

Timings on M3 Max MacBook Pro (16 CPU Cores, 40 GPU Cores)

With 100:
  In 1D:
    - cpu 0.004002000001491979
    - mps 0.10559125000145286
  In 2D:
    - cpu 0.0059719169803429395
    - mps 0.10134783299872652
  In 3D:
    - cpu 0.5764745000051335
    - mps 0.11991920799482614
With 250:
  In 1D:
    - cpu 0.001606041972991079
    - mps 0.08940245799021795
  In 2D:
    - cpu 0.026882333011599258
    - mps 0.09967358299763873
  In 3D:
    - cpu 6.992014332994586
    - mps 0.37927370800753124
With 500:
  In 1D:
    - cpu 0.0035721249878406525
    - mps 0.11156754201510921
  In 2D:
    - cpu 0.11158095800783485
    - mps 0.09940466698026285
  In 3D:
    - cpu 77.23693574999925
    - mps 3.229753541003447
With 1000:
  In 1D:
    - cpu 0.0015319999947678298
    - mps 0.6711677499988582
  In 2D:
    - cpu 0.29339966698898934
    - mps 0.10378754101111554
  In 3D:
    - cpu 428.23927774999174
    - mps 25.63869079199503