当完全访问/修改输出时，torch.conj_physical是否比torch.conj更快？

函数核心特性 #

torch.conj #

torch.conj() 执行延迟共轭操作，但可随时通过torch.resolve_conj()生成实际的共轭张量

torch.conj_physical #

无论共轭位是否设置，该函数都会执行共轭操作。

性能观察与测试 #

根据上述特性，当确定要访问/修改整个输出时，torch.conj_physical不会比torch.conj慢，有时甚至更快。我设计了如下测试函数：

def fn0(x):
    o = torch.conj(x)
    o += 1j

def fn1(x):
    o = torch.conj_physical(x)
    o += 1j

测试结果显示fn1明显更快，在CPU和GPU上速度提升5%-50%。但将第二行替换为torch.mean(o)或其他常见操作时，两者性能差异微乎其微（fn1仍略快）；替换为o *= x时仍能看到fn1的速度优势，这似乎与原地操作有关。

疑问 #

1. 当要访问/修改整个输出时，torch.conj是否存在比torch.conj_physical更快的情况？
2. 原地操作与非原地操作的性能差异原因是什么？

（环境：torch 2.0.1，Python 3.11.4，Windows 11）

完整基准测试脚本 #

# -*- coding: utf-8 -*-
import torch
import torch.utils.benchmark as benchmark

# Define funcs ###############################################################
def fn0(x):
    o = torch.conj(x)
    o += 1j

def fn1(x):
    o = torch.conj_physical(x)
    o += 1j

# Make bench funcs ###########################################################
names = ('fn0', 'fn1')
n_iters = 2000
n_iters_gpu = n_iters * 100
got_gpu = bool(torch.cuda.is_available())


#%% Benchmark ################################################################
times = {}
for N in (10000, 100000, 1000000):
    x = torch.randn(N, dtype=torch.complex64)
    x_gpu = x.cuda()

    times[N] = {}
    for name in names:
        common = dict(stmt=f'{name}(x)', setup=f'from __main__ import {name}')
        # cpu ----------------------------------------------------------------
        key = name
        bench_fn = benchmark.Timer(**common, globals={'x': x})
        # warmup
        _ = bench_fn.timeit(3)
        # bench
        times[N][key] = bench_fn.timeit(n_iters).mean

        # gpu ----------------------------------------------------------------
        if got_gpu:
            key = name + '-gpu'
            # warmup
            bench_fn = benchmark.Timer(**common, globals={'x': x_gpu})
            # warmup
            _ = bench_fn.timeit(300)
            # bench
            times[N][key] = bench_fn.timeit(n_iters_gpu).mean

        # "progress bar"
        print(end='.', flush=True)

#%% Print results ############################################################
print()
for N in times:
    print(f"N={N}")
    for name in names:
        print(name + '-cpu', "%.3g" % times[N][name])
    if got_gpu:
        for name in names:
            print(name + '-gpu', "%.3g" % times[N][name + '-gpu'])
    print()

内容的提问来源于stack exchange，提问作者OverLordGoldDragon