1 of 52

Опыт переноса cpu-bound задач дата-аналитики с Python на Rust

RustСon 2024

Кравченко В.С.

2 of 52

О себе

Python в области научных вычислений свыше 5 лет.

Rust ~ 2 года.

Интересуюсь численными методами и машинным обучением.

2

3 of 52

Benchmarks game

3

benchmarksgame-team.pages.debian.net

4 of 52

Energy Efficiency

4

Ranking Programming Languages by Energy Efficiency. Science of Computer Programming, volume 205. 2021

5 of 52

3 задачи

  1. Комбинаторная оптимизация (задача о множественном рюкзаке).�Декомпозиция временных рядов на временные ряды из заданного набора;
  2. Линейная регрессия. �Прогнозирование потребления природного газа (обучение+инференс);
  3. Задача линейного программирования. �Нахождение оптимального бронирования трубы для транспортировки газа.

5

6 of 52

Что мы сравниваем?

Python

  • CPython 3.12.7
  • Numpy 2.0.2�(OpenBLAS)
  • Scipy 1.14.1
  • Numba 0.60.0�(LLVM 14.0.6)

Python+Rust

  • CPython 3.13.7
  • Numpy 2.1.2
  • Nalgebra
  • Rust 1.82.0 �(LLVM 19.1.2)

6

Замеры проводились на AMD Ryzen 9 7590HS 64Gb

7 of 52

PyO3

7

8 of 52

PyO3. Стартуем Python расширение

my-project

├── Cargo.toml

├── my_project

│ ├── __init__.py

│ └── bar.py

├── pyproject.toml

├── README.md

└── src

└── lib.rs

8

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

#!/bin/bash

​

# Установка виртуального окружения для maturin

python -m venv .env

source .env/bin/activate

pip install maturin

​

# Запуск диалога создания проекта

maturin new

​

# Компиляция и установка в debug режиме

maturin develop

​

# Компиляция release версии и создание wheel

maturin build --release

9 of 52

PyO3. Hello World!

9

1

2

3

4

5

6

7

8

9

10

11

12

// lib.rs

use pyo3::prelude::*;

​

#[pymodule]

mod my_extension {

use super::*;

​

#[pyfunction]

fn sum(x: Vec<f64>) -> PyResult<f64> {

Ok(x.iter().sum())

}

}

1

2

3

4

// test.py

import my_extension

​

s = my_extension.sum([1.2, 2, 3, 5])

10 of 52

PyO3. Конвертация типов

10

1

2

3

4

5

6

7

8

9

10

11

12

13

14

// для #[pyfunction] или #[pymethods]

pub trait FromPyObject<'py>: Sized {

fn extract(ob: &'py PyAny) -> PyResult<Self> {

Self::extract_bound(&ob.as_borrowed())

}

// возможно устареет в 0.23

fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self>;

...

}

// для возвращаемых значений

pub trait IntoPy<T>: Sized {

fn into_py(self, py: Python<'_>) -> T;

...

}

11 of 52

PyO3. Передача аргументов

  1. Py<T>
  2. Borrowed<'a, 'py, T>
  3. Bound<'py, T>

​

11

PyAny

PyString

PyBool

PyFloat

PyList

PyDict

PyTuple

PySet

…

T

fn sum(x: Bound<'_, PyList>) -> Bound<'_, PyFloat>

12 of 52

PyO3. rust-numpy

pub struct PyReadonlyArray<'py, T, D>

where

T: Element,

D: Dimension,

{

array: Bound<'py, PyArray<T, D>>,

}

​

pub struct PyReadwriteArray<'py, T, D>

12

fn get_evens(numbers: PyReadonlyArray1<i64>) -> PyReadonlyArray1<i64>

Аллокация в куче Rust: IntoPyArray

Аллокация в куче NumPy: ToPyArray

13 of 52

Задача 1. Комбинаторная оптимизация

13

1

0 1 0 1 0

0 0 1 1 1

1 1 1 0 1

1 0 1 0 0

0 0 0 1 0

2

N

1

2

M

N = 215

M = 19

L = 1100

Набор 1

Комбинируемые

ряды

Набор 2

Целевые

ряды

…

…

1

2

1

2

N

M

…

…

1

5

1

3

8

=

9

2

=

MAPE

MAPE

+

+

-

14 of 52

Комбинаторный поиск. Результаты

14

15 of 52

Дата параллелизм

15

Распределённые вычисления

Многопоточность

SIMD

ILP

Timely Dataflow

Rayon

iterator_ilp

std::arch

std::simd

16 of 52

Многопоточность. Rayon

  • Data Parallelism
  • Work Stealing
  • Fork-Join Model

16

1

2

3

4

5

6

7

8

9

10

11

12

13

let numbers: Vec<i32> = (1..10000).collect();

let sum: i32 = numbers.par_iter().sum();

​

​

let data: Vec<i32> = (1..10000).collect();

data.par_chunks(100).for_each(|chunk| {

let result: i32 = chunk.iter().sum();

});

​

​

let (res1, res2) = rayon::join(

|| compute_task_1(),

|| compute_task_2());

17 of 52

Многопоточность. Результаты

17

18 of 52

SIMD (Single Instruction, Multiple Data)

// Скалярное тело цикла

vaddsd xmm0, xmm0, qword ptr [rdx + i]

vmovsd xmm1, qword ptr [rdi + i]

vmulsd xmm1, xmm1, qword ptr [rsi + i]

vaddsd xmm0, xmm0, xmm1

// Повторить код выше ещё 8 раза…

// Векторизованное тело цикла

vmovupd zmm0, zmmword ptr [rdi]

vmovupd zmm1, zmmword ptr [rsi]

vfmadd213pd zmm1, zmm0, zmmword ptr [rdx]

// Сложить все элементы в zmm1 и готово

18

a2

a1

x1

b1

y1

a3

a4

a5

a6

a7

a8

x2

x3

x4

x5

x6

x7

x8

b2

b3

b4

b5

b6

b7

b8

y2

y3

y4

y5

y6

y7

y8

y

∑

19 of 52

SIMD. Автовекторизация

19

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

// Этот код будет векторизован

fn sum_int(x: &[i64]) -> i64 {

x.iter().sum()

}

​

// А этот нет

fn sum_float(x: &[f64]) -> f64 {

x.iter().sum()

}

​

// И этот тоже нет

fn dot_float(a: &[f64], b: &[f64]) -> f64 {

assert!(a.len() >= 8);

assert_eq!(a.len(), b.len());

a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()

}

20 of 52

SIMD. std::arch

20

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

fn simd_sum(x: &[f64]) -> f64 {

let len = x.len();

let chunks = len / 8;

let mut sum = 0.0;

let mut sum_arr = [0.0f64; 8];

unsafe {

let mut sum_vec = _mm512_setzero_pd();

for i in 0..chunks {

let chunk_ptr = x.as_ptr().add(i * 8);

let vec = _mm512_loadu_pd(chunk_ptr);

sum_vec = _mm512_add_pd(sum_vec, vec);

}

_mm512_storeu_pd(sum_arr.as_mut_ptr(), sum_vec);

sum += sum_arr.iter().sum::<f64>();

for i in chunks * 8..len {

sum += *x.get_unchecked(i);

}

}

sum

}

// Это будет ошибкой

let vec =_mm512_load_pd(chunk_ptr);

21 of 52

SIMD. std::SIMD

pub struct Simd<T, const N: usize>([T; N])

where

LaneCount<N>: SupportedLaneCount,

T: SimdElement;

​

Simd<T,N> <=> [T; N]

21

1

2

3

4

5

6

7

8

const LANES: usize = 16;

fn simd_sum(x: &[f64]) -> f64 {

let (prefix, main, suffix) = x.as_simd::<LANES>();

let prefix_sum: f64 = prefix.iter().sum();

let suffix_sum: f64 = suffix.iter().sum();

let main_sum: Simd<f64, LANES> = main.iter().sum();

prefix_sum + suffix_sum + main_sum.reduce_sum()

}

22 of 52

ILP (Instruction-level parallelism)

22

Fetch

Decode

Execute

Write

1

Fetch

Decode

Execute

Write

2

3

4

5

6

7

8

Execute

Address Calculation

ALU Operations

Memory Access

Fetch

Decode

Execute

Write

9

A

B

C = f(A)

Cycle

23 of 52

Reciprocal throughput

VADDPD ~ 0.5 … 1

VMULPD ~ 0.5 … 1

​

https://www.agner.org/optimize/

https://uops.info/table.html

23

24 of 52

ILP. iterator_ilp

24

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

use iterator_ilp::IteratorILP;

const LANES: usize = 8;

fn simd_ilp_sum(x: &[f64]) -> f64 {

let (prefix, main, suffix) =

x.as_simd::<LANES>();

let prefix_sum = prefix.iter().sum();

let suffix_sum = suffix.iter().sum();

let main_sum = main

.iter()

.fold_ilp::<2, Simd<f64, LANES>>(

|| Simd::<f64, LANES>::splat(0.0),

|acc, &val| acc + val,

|acc, val| acc + val,

).reduce_sum();

prefix_sum + suffix_sum + main_sum

}

1

2

3

4

5

6

7

8

const LANES: usize = 16;

fn simd_sum(x: &[f64]) -> f64 {

let (prefix, main, suffix) = x.as_simd::<LANES>();

let prefix_sum: f64 = prefix.iter().sum();

let suffix_sum: f64 = suffix.iter().sum();

let main_sum: Simd<f64, LANES> = main.iter().sum();

prefix_sum + suffix_sum + main_sum.reduce_sum()

}

25 of 52

SIMD+ILP. Результаты

25

26 of 52

Simd. Общие рекомендации

  • Сначала разработайте не SIMD алгоритм
  • К ILP можно относится как расширению SIMD регистра
  • Проще расширить данные до кратного SIMD регистру размера
  • SIMD без оптимизации по кеш-локальности малоэффективен
  • Не забывайте про: �RUSTFLAGS=”-C target-feature=+avx2,+avx512f,+fma -C target-cpu=native" cargo build –release

26

27 of 52

PGO (Profile-guided optimization)

Что оптимизирует PGO?

  • Inlining
  • Branch Prediction
  • Code Layout
  • Function Layout
  • Value Profiling
  • Loop Optimizations
  • Indirect Call Promotion

​

27

28 of 52

Сборка с PGO

28

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

#!/bin/bash

​

# 1. Создаём инструментальную сборку библиотеки

export RUSTFLAGS="-Cprofile-generate=/tmp/pgo-data"

maturin develop --release

rm -rf /tmp/pgo-data

​

# 2. Снимаем профиль

python bench.py

​

# 3. Если запусков было несколько, то собираем их в один профиль

llvm-profdata merge -o /tmp/pgo-data/merged.profdata /tmp/pgo-data

​

# 4. Собираем библиотеку с использованием профиля

export RUSTFLAGS="-Cprofile-use=/tmp/pgo-data/merged.profdata"

maturin develop --release

Optimizing Rust programs with PGO and BOLT using cargo-pgo

29 of 52

PGO. Результаты

29

30 of 52

PGO. Общие рекомендации

  • Бенчмаркайте код до и после PGO.
  • Используйте стат. критерий для оценки ускорения (Манна–Уитни)
  • Возможен эффект “переоптимизации”

30

31 of 52

Линейная регрессия. Задача

  • Есть исторические ряды
    • Потребление газа
    • Средняя дневная температура, влажности, давление, световой день итд.
    • Календарные события: выходные, праздники, локдаун итд.
  • Нужно обучиться и делать прогнозы потребления газа, детектировать и фильтровать выбросы в данных

31

32 of 52

Линейная регрессия. Доп. требования

  • Работа в онлайн режиме
    • Ограничение на значения коэффициентов при обучени
    • Задание коэффициентов перед прогнозированием/сэмплированием
    • Фильтрация сэмплов по предикату
  • Интероп с:
    • Python (PyO3)
    • Julia (jlrs)
    • R (C ABI)
    • итд.

32

33 of 52

Линейная регрессия. Решение

33

сэмплы

Предикат 1

Предикат 2

Kfilter*

nalgebra

statsmodels

scipy

80 рядов длиной 1800, 10000 cемплов

* – модифицированный

VARMAX

+MGARCH

OLS

+GARCH

34 of 52

Линейная регрессия. Обучение

Обучение в 3 раза быстрее на Rust, но это не точно.

34

35 of 52

Линейная регрессия. Результаты

35

36 of 52

Задача 3. Линейное программирование

  • Есть набор прогнозных профилей объемов поставок газа
  • Есть цены на бронирование мощностей трубопровода.
    • Годовое бронирование 1
    • Бронирование на квартал 4
    • Бронирование на месяц 12
    • Бронирование на сутки вперёд (штраф) ~365

36

Объем поставки

Время

37 of 52

Линейное программирование

37

good_lp

Подготовка данных (Rust)

Результат

scipy

Подготовка данных (Python)

Результат

HiGHS

38 of 52

LP. Результаты

38

39 of 52

LP. Unsafe API

39

40 of 52

Субъективные ощущения

  • Переписывание с Python на Rust как оно есть не дает драматического прироста в производительности
  • Время затраченное на переписывание с Python - до 5x
  • Ошибки Rust в compile time дружественнее чем Python+Numba рантайма
  • Я смог перевести все расширения на Python 3.13 в день релиза
  • Нет окружения - нет проблем

40

41 of 52

Спасибо за внимание

41

​

​

Блоги по теме:

pythonspeed.com

mcyoung.xyz

pwy.io

terencezl.github.io/blog

thedataquarry.com/posts

Оцените доклад

42 of 52

Приложение

То, что не попало в основную часть,

отбракованные слайды

42

43 of 52

Что уже есть?

  • uv (pip)
  • Ruff
  • Pydantic 2.0
  • Polars (Pandas)
  • Lance (SQLite для эмбендингов)
  • orjson (быстрый json)
  • granian (быстрее Uvicorn)

43

44 of 52

SIMD. Автовекторизация

  • Используйте простые и понятные циклы без сложного потока управления.
  • Убедитесь, что шаблоны доступа к данным предсказуемы и желательно непрерывны.
  • Избегайте вызовов функций внутри цикла, которые нельзя встроить.
  • Минимизируйте зависимости между итерациями цикла.

44

45 of 52

scipy.signal.lfilter

a[0]*y[n] = b[0]*x[n] + b[1]*x[n-1] + ... + b[M]*x[n-M] - a[1]*y[n-1] - ... - a[N]*y[n-N]

​

45

lfilter(b, a, x, axis=-1, zi=None)

46 of 52

lfilter на std::Simd

46

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

impl<T, const N: usize> SimdExt<T, N> for Simd<T, N>

where

T: SimdElement,

Simd<T, N>: Add<Output = Simd<T, N>> + Mul<Output = Simd<T, N>>,

LaneCount<N>: SupportedLaneCount,

Simd<T, N>: SimdFloat<Scalar = T>,

{

fn push(&mut self, item: T) {

*self = self.rotate_elements_left::<1>();

self[N - 1] = item;

}

​

fn dot(&self, vals: &Simd<T, N>) -> T {

(*self * *vals).reduce_sum()

}

}

47 of 52

lfilter на std::Simd

47

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

pub fn lfilter<const N: usize, const M: usize>(

a: &[f64; N], b: &[f64; M],

x: &[f64], zi: &[f64; N],

) -> Vec<f64>

where

LaneCount<N>: SupportedLaneCount, LaneCount<M>: SupportedLaneCount,

{

let a_coefs: Simd<_, M> = Simd::from_slice(b);

let b_coefs: Simd<_, N> = Simd::from_slice(a);

let mut a_part = Simd::<_, M>::splat(0.0);

let mut b_part = Simd::<_, N>::from_slice(zi);

let mut result = vec![0.0; x.len()];

for (i, y) in result.iter_mut().enumerate() {

a_part.push(*y);

b_part.push(x[i]);

*y = b_part.dot(&b_coefs) + a_part.dot(&a_coefs);

}

result

}

Ускорение:

  • 2X по сравнению с наивным Rust
  • ~2.2X по сравнению с lfilter из Scipy

48 of 52

SIMD. Что за зверь?

vmovupd zmm0, zmmword ptr [rdi]

vmovupd zmm1, zmmword ptr [rsi]

vfmadd213pd zmm1, zmm0, zmmword ptr [rdx]

// Сложить все все элементы в zmm1

// y += a[i] * b[i] + c[i]

vmovupd xmm0, xmmword ptr [rdi]

vmulpd xmm0, xmm0, xmmword ptr [rsi]

vaddpd xmm0, xmm0, xmmword ptr [rdx]

vaddsd xmm1, xmm0, xmm1

vshufpd xmm0, xmm0, xmm0, 1

vaddsd xmm0, xmm1, xmm0

// Повторить ещё 4 раза…

// y = y + a[i] * b[i] + c[i]

vaddsd xmm0, xmm0, qword ptr [rdx + i]

vmovsd xmm1, qword ptr [rdi + i]

vmulsd xmm1, xmm1, qword ptr [rsi + i]

vaddsd xmm0, xmm0, xmm1

// Повторить ещё 8 раза…

48

a2

a1

x1

b1

y1

a3

a4

a5

a6

a7

a8

x2

x3

x4

x5

x6

x7

x8

b2

b3

b4

b5

b6

b7

b8

y2

y3

y4

y5

y6

y7

y8

y

∑

49 of 52

ILP (Instruction-level parallelism)

IF (англ. instruction fetch) — узел, ответственный за получение инструкции;

ID (англ. instruction decode) — узел, ответственный за раскодирование инструкции;

EX (англ. execute) — узел, ответственный за выполнение инструкции;

MEM (англ. memory access) — узел, ответственный за доступ к памяти;

WB (англ. register write back) — узел, ответственный за запись в регистр.

​

49

50 of 52

SIMD. -ffast-math

cargo-features = ["profile-rustflags"]

rustflags = [

"-C", "target-feature=+avx2,+avx512f,+fma",

"-C", "llvm-args=-ffast-math",

"-C", "llvm-args=-force-vector-width=8",

]

​

https://llvm.org/docs/Vectorizers.html

50

51 of 52

SIMD. std::SIMD

pub struct Simd<T, const N: usize>([T; N])

where

LaneCount<N>: SupportedLaneCount,

T: SimdElement;

​

  • по большей части Safe
  • высокоуровневая
  • experimental

​

Simd<T,N> <=> [T; N]

​

51

52 of 52

HiGHS

52