Несколько лет назад я написал реализацию алгоритма умножения матриц Штрассена на C++, а недавно повторно реализовал его на Rust, продолжая изучать язык. Это было полезное упражнение для изучения характеристик производительности и методов оптимизации Rust, потому что, хотя алгоритмическая сложность Штрассена превосходит наивный подход, он имеет высокий постоянный коэффициент от накладные расходы на выделение и рекурсию в структуре алгоритма.

Общий алгоритм

Общий (наивный) алгоритм умножения матриц — это подход с тремя вложенными циклами, который каждый изучает на своем первом уроке линейной алгебры, который большинство распознает как O(n³).

pub fn 
mult_naive (a: &Matrix, b: &Matrix) -> Matrix {
    if a.rows == b.cols {
        let m = a.rows;
        let n = a.cols;

        // preallocate
        let mut c: Vec<f64> = Vec::with_capacity(m * m);

        for i in 0..m {
            for j in 0..m {
                let mut sum: f64 = 0.0;
                for k in 0..n {
                    sum += a.at(i, k) * b.at(k, j);
                }

                c.push(sum);
            }
        }

        return Matrix::with_vector(c, m, m);
    } else {
        panic!("Matrix sizes do not match");
    }
}

Этот алгоритм работает медленно не только из-за трех вложенных циклов, но и потому, что внутренний цикл обхода B по столбцам через b.at(k, j), а не по строкам, ужасен для скорости попадания в кэш ЦП.

Транспонирование для лучшей производительности

Транспонированный наивный подход реорганизует шаги умножения матрицы B в более благоприятный для кэширования формат, позволяя итерациям умножения над B выполняться по строкам, а не по столбцам. Таким образом A x B превращается в A x B^t

Он включает в себя новое распределение матрицы (во всяком случае, в этой реализации) и полную итерацию матрицы (операция O(n²), то есть, точнее, этот подход O(n³) + O (n²)) — ниже я покажу, насколько лучше он работает. Это выглядит следующим образом:

fn multiply_transpose (A: Matrix, B: Matrix):
  C = new Matrix(A.num_rows, B.num_cols)

  // Construct transpose; requires allocation and iteration through B
  B’ = B.transpose()

  for i in 0 to A.num_rows:
    for j in 0 to B'.num_rows:
      sum = 0;
      for k in 0 to A.num_cols:
        // Sequential access of B'[j, k] is much faster than B[k, j]
        sum += A[i, k] * B'[j, k]
      C[i, j] = sum
  return C

Субкубический: как работает алгоритм Штрассена

Чтобы понять, как работает алгоритм Штрассена (код на Rust здесь), сначала рассмотрим, как матрицы могут быть представлены квадрантами. Чтобы понять, как это выглядит:

Используя эту квадрантную модель в наивном алгоритме, каждый из четырех квадрантов результирующей матрицы C представляет собой сумму произведений двух подматриц, что дает всего 8 умножений.

Принимая во внимание эти восемь умножений, каждое из которых работает с блочной матрицей с размахом строк и столбцов примерно в половину размера A и B, сложность получается такой же:

Алгоритм Штрассена определяет семь матриц промежуточных блоков, составленных из следующих квадрантов:

Которые вычисляются с помощью только 7 умножений, а не 8. Эти умножения могут быть рекурсивными умножениями Штрассена и могут использоваться для составления окончательной матрицы как таковой:

С полученной субкубической сложностью:

Параллелизм

Вычисление промежуточных матриц M1 … M7 является досадно параллельной задачей, поэтому было легко инструментировать и параллельный вариант алгоритма (как только вы начнете понимать правила Rust относительно замыканий).

/**
 * Execute a recursive strassen multiplication of the given vectors, 
 * from a thread contained within the provided thread pool.
 */
fn 
_par_run_strassen (a: Vec<f64>, b: Vec<f64>, 
                   m: usize, pool: &ThreadPool) 
                     -> Arc<Mutex<Option<Matrix>>> {
    let m1: Arc<Mutex<Option<Matrix>>> = Arc::new(Mutex::new(None));
    let m1_clone = Arc::clone(&m1);
     
    pool.execute(move|| { 
        // Recurse with non-parallel algorithm once we're 
        // in a working thread
        let result = mult_strassen(
            &mut Matrix::with_vector(a, m, m),
            &mut Matrix::with_vector(b, m, m)
        );
        
        *m1_clone.lock().unwrap() = Some(result);
    });

    return m1;
}

Бенчмаркинг

Я написал код быстрого бенчмаркинга, который запускает каждый из четырех алгоритмов в течение нескольких испытаний в увеличивающемся диапазоне размеров матрицы и сообщает среднее время для каждого из них.

~/code/strassen ~>> ./strassen --lower 75 --upper 100 --factor 50 --trials 2

running 50 groups of 2 trials with bounds between [75->3750, 100->5000]

x    y    nxn      naive       transpose  strassen   par_strassen
75   100  7500     0.00ms      0.00ms     1.00ms     0.00ms
150  200  30000    6.50ms      4.00ms     4.00ms     1.00ms
225  300  67500    12.50ms     9.00ms     8.50ms     2.50ms
300  400  120000   26.50ms     22.00ms    18.00ms    5.50ms
[...]
3600 4800 17280000 131445.00ms 53683.50ms 21210.50ms 5660.00ms
3675 4900 18007500 141419.00ms 58530.00ms 28291.50ms 6811.00ms
3750 5000 18750000 154941.00ms 60990.00ms 26132.00ms 6613.00ms

Затем я визуализировал результаты через pyplot :

На этом графике показано время умножения матриц от 7,5 тыс. элементов (N x M = 75 x 100) до примерно 19 миллионов (N x M = 3750 x 5000). Вы можете видеть, как быстро наивный алгоритм становится непрактичным с вычислительной точки зрения, занимая две с половиной минуты на верхнем уровне.

Алгоритм Штрассена, для сравнения, масштабируется более плавно, а параллельный алгоритм вычисляет результаты для двух матриц по 19 миллионов элементов в каждой за то время, которое потребовалось наивному алгоритму для обработки всего 3,6 миллиона элементов.

Больше всего меня интересует производительность алгоритма transpose. Как обсуждалось ранее, улучшение производительности кэша за счет полного копирования матрицы ясно продемонстрировано в этих результатах — даже с алгоритмом, асимптотически эквивалентным подходу naive.

Профилирование и оптимизация производительности

Эта документация была фантастическим ресурсом для понимания основ производительности Rust. Настроить и запустить Инструменты для профилирования в Mac OS было тривиально благодаря Руководству по грузовым инструментам Rust. Это отличный инструмент как для исследования поведения распределения, горячих точек процессора, так и для других вещей.

Некоторые изменения по ходу:

  • Код Штрассена вызывает себя рекурсивно с помощью стратегии разделяй и властвуй, но как только матрицы достигают достаточно малого размера, его высокий постоянный коэффициент делает его медленнее, чем общий матричный алгоритм. Я обнаружил, что эта точка представляет собой строку или столбец шириной около 64; увеличение этого порога с 2 увеличило пропускную способность на несколько факторов
  • Алгоритм Штрассена требует заполнения матрицы до ближайшего показателя степени 2; уменьшение этого, чтобы лениво гарантировать, что матрицы имеют только четное количество строк и столбцов, повысило пропускную способность примерно вдвое за счет сокращения дорогостоящих больших распределений.
  • Изменение алгоритма отката для малой матрицы с naive на transpose привело к улучшению примерно на 20%.
  • Добавление codegen-units = 1 и lto = "thin" к флагам сборки релиза Cargo.toml дало улучшение примерно на 5%. Интересно, что lto = “true” постоянно ухудшал производительность.
  • Тщательное удаление всех возможных Vec копий дало улучшение примерно на 10%.
  • Предоставление пары#[inline] подсказок и удаление проверки границ вектора при поиске с произвольным доступом привели к улучшению еще примерно на 20%.
    /**
     * Returns the element at (i, j). Unsafe.
     */
    #[inline]
    pub fn at (&self, i: usize, j: usize) -> f64 {
         unsafe {
            return *self.elements.get_unchecked(i * self.cols + j);
        }
    }