RTX4090のメモリ転送の謎を考察する

# -*- coding: utf-8 -*-
import numpy as np
import torch
import datetime

def diff_time(st, ed):
    d = ed - st
    return d.seconds + (d.microseconds / 1000000.)

def test_body(title, a, a_rhs):
    st = datetime.datetime.now()
    for i in range(lattice):
        if i == lattice - 1:
            st_slice = datetime.datetime.now()
        a_rhs[:, :] = a[:, :]
        if i == lattice - 1:
            ed_slice = datetime.datetime.now()
            print('a to a_rhs:', diff_time(st_slice, ed_slice))
        if i == lattice - 1:
            st_slice = datetime.datetime.now()
        a[:, :] = a_rhs[:, :]
        if i == lattice - 1:
            ed_slice = datetime.datetime.now()
            print('a_rhs to a:', diff_time(st_slice, ed_slice))
    ed = datetime.datetime.now()
    print(title, diff_time(st, ed))

def test_round(lattice, dev):
    print("lattice:{}({})".format(lattice, dev))
    x = np.random.rand(lattice, lattice)
    x_64 = torch.from_numpy(x).to(dtype=torch.float64)
    x_32 = torch.from_numpy(x).to(dtype=torch.float32)
    x_rhs = np.zeros((lattice, lattice), dtype=np.float64)

    test_body('cpu:', x, x_rhs)
    x_rhs = torch.zeros((lattice, lattice), dtype=torch.float32).to(device=dev)
    test_body('gpu(32):', x_32, x_rhs)
    x_rhs = torch.zeros((lattice, lattice), dtype=torch.float64).to(device=dev)
    test_body('gpu(64):', x_64, x_rhs)
    print()

if __name__ == '__main__':
    lattices = [2400, 2400]
    devs = ['cpu', 'cuda:0']
    for lattice, dev in zip(lattices, devs):
        test_round(lattice, dev)

lattice:2400(cpu)
a to a_rhs: 0.002171
a_rhs to a: 0.002289
cpu: 10.85354
a to a_rhs: 0.000217
a_rhs to a: 0.000209
gpu(32): 1.071928
a to a_rhs: 0.001655
a_rhs to a: 0.001987
gpu(64): 9.114049

lattice:2400(cuda:0)
a to a_rhs: 0.002173
a_rhs to a: 0.002306
cpu: 10.782213
a to a_rhs: 0.001854
a_rhs to a: 0.001967
gpu(32): 9.739486
a to a_rhs: 0.003724
a_rhs to a: 0.004149
gpu(64): 18.750687