Numpy einsum在特定问题规模后性能骤降原因排查

运行该脚本仅需依赖numpy和matplotlib。

import numpy as np
import timeit
import matplotlib.pyplot as plt

n1, n2, n3, n4, n5, m = (101, 33, 1, 32, 2, 32)


def make_arrays(aOrder, bOrder, cOrder):
    a = np.random.randn(n1, m) + 1j * np.random.randn(n1, m)
    b = np.random.randn(n2, m) + 1j * np.random.randn(n2, m)
    c = np.random.randn(n1, n2, n3, n4, n5) + 1j * np.random.randn(n1, n2, n3, n4, n5)
    return (
        np.array(a, order=aOrder),
        np.array(b, order=bOrder),
        np.array(c, order=cOrder),
    )


# used in B()
blockSize = 84

resA = []
resB = []
resC = []

sizes = np.unique(np.exp(np.linspace(2, 6, 8)).astype(np.int64))
numTrials = 10

# overrides m form above
for m in sizes:
    a, b, c = make_arrays("F", "F", "F")

    path = np.einsum_path(
        a,
        [0, 5],
        b,
        [1, 5],
        c,
        [0, 1, Ellipsis],
        [Ellipsis, 5],
        optimize="greedy",
    )

    def A():
        np.einsum(
            a,
            [0, 5],
            b,
            [1, 5],
            c,
            [0, 1, 2, 3, 4],
            [5, 2, 3, 4],
            optimize="greedy",
            order="F",
        )
        # print("einsum\n", res.flags)

    def B():
        numBlocks = int(a.shape[1] // blockSize)
        np.concatenate(
            tuple(
                np.einsum(
                    c,
                    [1, 2, Ellipsis],
                    a[:, kk * blockSize : (kk + 1) * blockSize],
                    [1, 0],
                    b[:, kk * blockSize : (kk + 1) * blockSize],
                    [2, 0],
                    [0, Ellipsis],
                    optimize="greedy",
                    order="F",
                )
                for kk in range(numBlocks)
            )
            + (
                np.einsum(
                    c,
                    [1, 2, Ellipsis],
                    a[:, numBlocks * blockSize :],
                    [1, 0],
                    b[:, numBlocks * blockSize :],
                    [2, 0],
                    [0, Ellipsis],
                    optimize="greedy",
                    order="F",
                ),
            ),
            axis=0,
        )

    def C():
        tmp = np.einsum(a, [0, 5], c, [0, 1, 2, 3, 4], [1, 2, 3, 4, 5], order="F")
        np.einsum(b, [1, 5], tmp, [1, 2, 3, 4, 5], [2, 3, 4, 5], order="F")

    A()
    B()
    C()

    measured = np.mean(timeit.repeat(A, number=numTrials, repeat=numTrials)) / (
        numTrials * m
    )
    resA.append(measured)

    measured = np.mean(timeit.repeat(B, number=numTrials, repeat=numTrials)) / (
        numTrials * m
    )
    resB.append(measured)

    measured = np.median(timeit.repeat(C, number=numTrials, repeat=numTrials)) / (
        numTrials * m
    )
    resC.append(measured)

plt.figure()
plt.semilogy(sizes, resA, label="A")
plt.semilogy(sizes, resB, label="B")
plt.semilogy(sizes, resC, label="C")

plt.legend()
plt.show()

方法说明 #

• 方法A：使用einsum的常规方式，原本预期性能表现良好，但当问题规模m超过某个阈值后，计算速度会骤降。
• 方法B：分块实现，全规模下性能表现稳定，但在m较小时性能弱于A。无论如何调整输入和输出数组的内存布局，该定性表现均无明显变化。
• 方法C：更为基础的分步实现，性能极差，符合预期。

测试环境与结果 #

AMD EPYC环境 #

配置：搭载AMD EPYC 7343 16核处理器、使用IntelPython（基于MKL）且强制MKL仅使用单核心
说明：将运行时间除以问题规模m，以此估算单个问题“切片”的计算成本

M2 MacBook Air环境 #

原生numpy测试

配置：搭载M2芯片的MacBook Air，使用原生numpy

链接Accelerate BLAS库的numpy测试

配置：搭载M2芯片的MacBook Air，使用链接Accelerate BLAS库的numpy

核心问题 #

方法A到底出了什么问题？

题外话 #

• 在CuPy上运行相同代码，未出现该现象，对应A方法的实现可良好扩展，至少在上述相同问题规模下表现正常。
• 若使用opt_einsum进行收缩运算（基本沿用A的代码逻辑），可得到类似B的性能，但稍慢一些。

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