행렬 연산

제로 카피 연산에서 단순히 더하고 곱하는 건 메모리 대역폭 싸움이었다면,
행렬 곱셈은 CPU 연산 능력과 데이터 재사용의 싸움이다.

64GB 램 안에서 가장 거대하고 빠른 행렬 연산을 구현해 보도록 하자.

작업공간 생성 및 구조 확인

~/workspace$ mkdir matrix-multiply && cd matrix-multiply
~/workspace/matrix-multiply$ python3 -m venv venv
~/workspace/matrix-multiply$ source venv/bin/activate
~/workspace/matrix-multiply$ # 평소에 설치하던 것 외에 추가된 라이브러리를 놓치지 말자
~/workspace/matrix-multiply$ pip install maturin fastapi uvicorn orjson numpy
~/workspace/matrix-multiply$ maturin init
~/workspace/matrix-multiply$ # 선택지 중 기본값인 PyO3 선택
~/workspace/matrix-multiply$ # Cargo.toml과 src/lib.rs가 자동 생성된다
~/workspace/matrix-multiply$ # Python 코드는 직접 생성해 주어야 한다
~/workspace/matrix-multiply$ mkdir app && touch app/main.py
~/workspace/matrix-multiply$ tree -I venv
.
├── app
│   └── main.py
├── Cargo.toml
├── pyproject.toml
└── src
    └── lib.rs

Cargo.toml 파일을 열어 라이브러리 이름을 수정해 주겠다.

이번 실습에서는 numpy 크레이트와 함께
ndarray 크레이트를 사용한다고 명시해 둔다.
numpy 크레이트가 내부적으로 ndarray 크레이트를 가지고 사용하고 있어도
Rust의 의존성 관리 규칙에 의해
다음과 같은 이유로 명시적으로 작성해주는 것을 권장한다.

Cargo.toml

[package]
name = "matrix-multiply"
version = "0.1.0"
edition = "2024"

[lib]
name = "rust_engine"
crate-type = ["cdylib"]

[dependencies]
pyo3 = "0.28.0"
numpy = "0.28"
ndarray = "0.17"

코드 작성

Rust 코드

지금까지는 1차원 배열을 사용했지만
이번에는 행렬 연산을 수행할 것이기에 2차원 배열을 가져온다.

src/lib.rs

use pyo3::prelude::*;
use numpy::{PyReadonlyArray2, PyArray2};

#[pyfunction]
fn fast_matrix_multiply(
    py: Python<'_>,
    mat_a: PyReadonlyArray2<f64>,
    mat_b: PyReadonlyArray2<f64>
) -> Py<PyArray2<f64>> {
    let a = mat_a.as_array();
    let b = mat_b.as_array();

    let result = a.dot(&b);

    PyArray2::from_owned_array(py, result).unbind()
}

#[pymodule]
fn rust_engine(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fast_matrix_multiply, m)?)?;

    Ok(())
}

Python 코드

app/main.py

import numpy as np
from fastapi import FastAPI
from fastapi.responses import ORJSONResponse
import rust_engine
import time

class UTF8ORJSONResponse(ORJSONResponse):
    media_type = "application/json; charset=utf-8"

app = FastAPI(default_response_class=UTF8ORJSONResponse)

@app.get("/")
def read_root():
    return {
        "status": "200",
        "info": "서버 가동 중입니다."
    }

@app.get("/matrix/{dim}")
def matrix_test(dim: int):
    a = np.random.rand(dim, dim).astype(np.float64)
    b = np.random.rand(dim, dim).astype(np.float64)

    start = time.perf_counter()
    result = rust_engine.fast_matrix_multiply(a, b)
    end = time.perf_counter()

    rust_duration = end - start

    return {
        "dimension": f"{dim}x{dim}",
        "total_elements": dim*dim,
        "rust_pure_time": f"Rust 연산에 걸린 시간: {rust_duration:.4f} sec",
        "result_sample": result[0, :3].tolist()
    }

빌드 및 실행

Maturin 라이브러리를 통해 Rust 코드를 Python에서 호출 가능한 형태로 컴파일한다.
ndarray 크레이트의 dot() 연산이 내부적으로 병렬 처리를 지원하므로
성능 최적화를 위해 --release 를 붙여 컴파일한다.
컴파일 후 pip list 명령어를 사용해 보면 Cargo.toml 파일에 작성한 패키지 이름을 확인할 수 있다.

uvicorn 라이브러리를 통해 FastAPI를 실행한다.

~/workspace/matrix-multiply$ maturin develop --release
~/workspace/matrix-multiply$ uvicorn app.main:app --reload

curl 명령어 또는 브라우저를 통해 다음과 같은 테스트를 해볼 수 있다.

• http://127.0.0.1:8000/matrix/10000

~$ curl -i http://127.0.0.1:8000/matrix/10000
HTTP/1.1 200 OK
date: Mon, 30 Mar 2026 21:52:19 GMT
server: uvicorn
content-length: 190
content-type: application/json; charset=utf-8

{"dimension":"10000x10000","total_elements":100000000,"rust_pure_time":"Rust 연산에 걸린 시간: 37.1463 sec","result_sample":[2493.7760400667707,2484.6190913959954,2469.6911490233115]}%   

행렬 연산은 O(N³) 연산이므로 데이터 크기가 늘어날수록 세제곱에 비례하게 속도가 느려진다.
이 테스트는 약 2.4GB의 데이터를 다루며 1조 번의 곱셈과 덧셈을 수행한다.
데이터가 큰 만큼 오래 걸린 것 같아 보이지만 초당 약 538억 번의 부동 소수점 연산이 있었다.
Accelerate 프레임워크 같은 하드웨어 가속 라이브러리를 연결하여
하드웨어 최적화를 하면 속도를 훨씬 더 줄일 수 있겠지만 여기서는 다루지 않겠다.