Явное умножение матриц намного быстрее, чем numpy.matmul?

Проверка грубой линейности matmul по первому, «пакетному» измерению:

Ваши (1000,2,2) массивы:

In [353]: timeit matrices_a@matrices_b
251 µs ± 767 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

и с половинным и десятым размером:

In [354]: timeit matrices_a[:500]@matrices_b[:500]
129 µs ± 783 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [355]: timeit matrices_a[:100]@matrices_b[:100]
28.7 µs ± 532 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

Ваш явный

In [360]: explicit_2x2_matrices_multiplication(matrices_a, matrices_b).shape
Out[360]: (1000, 2, 2)
In [361]: timeit explicit_2x2_matrices_multiplication(matrices_a, matrices_b)
59.9 µs ± 1.91 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

np.einsum не пытается изменить порядок или другую оптимизацию:

In [362]: print(np.einsum_path('ijk,ikl->ijl',matrices_a, matrices_b, optimize='greedy
     ...: ')[1])
  Complete contraction:  ijk,ikl->ijl
         Naive scaling:  4
     Optimized scaling:  4
      Naive FLOP count:  1.600e+04
  Optimized FLOP count:  1.600e+04
   Theoretical speedup:  1.000
  Largest intermediate:  4.000e+03 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   4                ikl,ijk->ijl                                 ijl->ijl

TL; DR: все методы, представленные в вопросе, очень неэффективны. На самом деле, Numpy явно неоптимален, и нет никакого способа эффективно вычислить это только с помощью Numpy. Тем не менее, есть более быстрые решения, чем те, которые указаны в вопросе.

Объяснение и более быстрая реализация

Код Numpy использует мощные универсальные итераторы для применения заданного вычисления (например, умножения матриц) к нескольким срезам массива. Такие итераторы полезны для реализации вещания, а также для создания относительно простой реализации einsum. Дело в том, что они также довольно дороги, когда количество итераций огромно, а массив небольшой. Это именно то, что происходит в вашем случае использования. Хотя накладные расходы итераторов Numpy можно уменьшить за счет оптимизации кода Numpy, в этом конкретном случае использования невозможно сократить накладные расходы до незначительного времени. Действительно, на каждую матрицу приходится выполнять только 12 операций с плавающей запятой. Относительно новый процессор x86-64 может вычислять каждую матрицу менее чем за 10 наносекунд, используя скалярные единицы FMA. Фактически, можно использовать инструкцию SIMD, чтобы вычислить каждую матрицу всего за несколько наносекунд.

Во-первых, мы можем в основном устранить накладные расходы внутреннего итератора Numpy, самостоятельно выполняя умножение матриц, оперируя векторами вдоль первой оси. Это именно то, что делает explicit_2x2_matrices_multiplication!

Хотя explicit_2x2_matrices_multiplication должен быть значительно быстрее, он все еще неоптимален: он выполняет несмежное чтение памяти, создает несколько бесполезных временных массивов, и каждый вызов Numpy вносит небольшие накладные расходы. Более быстрое решение — написать код Numba/Cython, чтобы базовый компилятор мог генерировать очень хорошую последовательность инструкций, оптимизированную для матриц 2x2. В этом случае хорошие компиляторы могут даже генерировать SIMD-инструкции, что невозможно для Numpy. Вот полученный код:

import numba as nb
@nb.njit('(float64[:,:,::1], float64[:,:,::1])')
def compute_fastest(matrices_a, matrices_b):
    assert matrices_a.shape == matrices_b.shape
    assert matrices_a.shape[1] == 2 and matrices_a.shape[2] == 2

    n = matrices_a.shape[0]
    result_matrices = np.empty((n, 2, 2))
    for k in range(n):
        for i in range(2):
            for j in range(2):
                result_matrices[k,i,j] = matrices_a[k,i,0] * matrices_b[k,0,j] + matrices_a[k,i,1] * matrices_b[k,1,j]

    return result_matrices

Результаты производительности

Вот результаты производительности на моей машине с процессором i5-9600KF для матриц 1000x2x2:

Naive einsum:                           214   µs
matrices_a @ matrices_b:                102   µs
explicit_2x2_matrices_multiplication:    24   µs
compute_fastest:                          2.7 µs   <-----

Обсуждение

Реализация Numba достигает 4,5 GFlops. Каждая матрица вычисляется всего за 12 циклов процессора (2,7 нс)! Моя машина на практике может достигать 300 GFlops (теоретически 432 GFlops), но только 50 GFlops с 1 ядром и 12,5 GFlops при использовании скалярного кода (теоретически 18 GFlops). Детализация операции слишком мала, чтобы несколько потоков могли быть полезными (накладные расходы на порождающий поток составляют как минимум несколько микросекунд). Кроме того, код SIMD вряд ли может насытить единицы FMA, потому что структура входных данных требует перетасовки SIMD, поэтому 50 GFlops на самом деле является оптимистичной верхней границей. В результате мы можем с уверенностью сказать, что реализация Numba — довольно эффективная реализация. Тем не менее, благодаря SIMD-инструкциям можно написать более быстрый код (на практике я ожидаю ускорения примерно в 2 раза). При этом написание нативного кода с использованием встроенных функций SIMD, помогающих компилятору генерировать быстрый код SIMD, на самом деле далеко не просто (не говоря уже о том, что окончательный код будет уродливым и сложным в обслуживании). Таким образом, реализация SIMD может не стоить затраченных усилий.