Несколько лет назад я написал реализацию алгоритма умножения матриц Штрассена на 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); } }