Rust와 AI에서 메모리 아끼기: 고정폭 비트 패킹 정수 벡터의 매력

거대한 AI 그리고 인공지능 데이터셋을 다룰 때, 메모리 효율은 곧 성능입니다. Rust 언어로 고정폭(bit-width) 비트 패킹 정수 벡터를 직접 만드는 방법을 살펴보며, 어떻게 데이터 압축과 O(1) 속도의 빠른 접근을 한 번에 잡을 수 있는지 흥미진진하게 풀어봅니다.

메모리 낭비의 진실: 왜 Vec<T>만 써서는 안 될까?

Rust의 Vec<T>는 쉽고 빠른 O(1) 접근을 보장하죠. 하지만, 데이터를 정말 큰 배열에 저장할 때 문제가 시작됩니다! 예를 들어서, u64 타입(8바이트)로 저장되는 값이 실제로는 3비트만 있으면 표현되는 상황을 상상해보세요.
10억 개가 필요하다면? 단순히 8GB가 필요하지만, 진짜 실제 데이터는 1GB도 안 될 수도 있습니다. 이런 엄청난 낭비, 해결법이 없을까요?

정답은 비트 패킹! 필요한 만큼의 비트만 딱 맞춰서 데이터를 쭉쭉 이어붙이면, 필요한 메모리만큼만 써서 저장할 수 있습니다. 하지만 비트 단위로 데이터를 저장하면, 원하는 값을 바로 찾기가 어려워질 것 같다는 게 옛날 고민이었습니다.

O(1) 성능을 위한 비트 패킹 해법

비트 패킹의 핵심은 "내가 찾을 값이 어느 위치에 있지?"를 천재적으로 계산하는 겁니다.
예를 들어, max값이 1000이면 모든 데이터는 10비트 내에서 저장 가능하죠.
각 원소는 index * bit_width만큼 떨어져 있으니, 간단한 계산으로 어느 비트에 값이 있는지 알 수 있습니다.

몇 번째 값인지 index를 곧장 생각한다
필요한 비트만큼만 추출한다
마스킹(mask)으로 필요 없는 부분은 지운다

이렇게 하면 원하는 값을 O(1) 시간에 바로 찾을 수 있습니다.

비트 경계 넘기: 마법의 두 단어 비트 결합

하나의 u64(64비트) 워드에 값이 나눠 들어가는 상황, 예를 들어 10비트씩 저장한다면 어떤 값은 한 워드의 끝과 다음 워드까지 걸칠 수밖에 없습니다.

이런 상황에선 두 워드에서 해당 비트를 뽑아내어 합쳐준 뒤, 마스킹 작업을 하게 됩니다. 데이터가 워드 경계를 넘어가도 성능은 크게 저하되지 않아요―끝부분만 살짝 더 복잡할 뿐!

추가로, 마지막 워드는 안전하게 패딩해서, 마지막 값이 경계를 넘어가도 항상 유효한 메모리 위치를 가리키도록 해야 합니다.

CPU와 캐시를 이기는 비트 패킹 접근법

읽기 성능을 최대치로 끌어올리기 위해, 단일 비트에서 시작하는 데이터의 "첫 바이트 위치"를 계산해, CPU가 더 빠르게 읽을 수 있도록 합니다.
이때 Rust의 read_unaligned을 사용하면 CPU가 알아서 비트 경계를 건너뛰는 작업까지 한 번에 해줍니다.
더이상 두 번 데이터 읽고 합치는 복잡한 과정이 아니라, 한 번에 읽고, 바로 마스킹과 쉬프트로 결과를 가져옵니다.

결과적으로, 작은 비트폭(예: 4~32비트)는 L1캐시 적중률이 높아, 일반적인 Vec<u8>보다도 더 빠른 무작위 접근 성능을 보일 수 있습니다.

빠른 반복: 비트 패킹 벡터의 스마트 이터레이터

for문과 반복문에서 효율은 더욱 중요합니다.
비트 패킹 벡터는 반복문을 돌 때, 단순히 get()을 매번 호출하면 계산 반복이 많아집니다.
대신 현재 워드와 비트 위치를 내부적으로 기억하는 스테이트풀(stateful) 이터레이터를 만들어, 필요한 비트만 쭉쭉 뽑아오게 설계할 수 있습니다.

그리고 바로 양방향(앞/뒤) 이터레이션까지 지원!
앞쪽에선 front_window, 뒤쪽에선 back_window와 비트 카운터를 유지하면서, 원하는 방향으로 더욱 빠르게 반복이 가능합니다.

비트 단위 쓰기: 안전하고 빠른 데이터 수정

읽기만 빠르면 아쉬우니, 집어넣는 것도 생각해야죠!
비트 경계 내에 값이 있는 경우엔 해당 비트 범위만 깨끗하게 지우고, 새 값을 합쳐 넣습니다.
경계를 넘어가면? 두 워드를 각각 부분적으로 수정해서 데이터 일관성을 깔끔하게 유지합니다.
이때 Rust의 안전한 split_at_mut_unchecked 활용하면 두 메모리 위치를 중복 걱정 없이 효율적으로 접근할 수 있습니다.

실제 무작위 쓰기 성능은 일반 Vec<T>보다 살짝 느릴 수 있지만, 공간 절약과 읽기 성능의 이점이 훨씬 큽니다.

제네릭 구조와 확장성: 모든 상황을 위한 유연한 설계

단순 u64 전용 타입이 아니기에, 구조체와 trait를 활용해 로직 타입, 물리적 저장(word), 엔디안, 저장 방식까지 원하는 대로 조합 가능하게 설계합니다.

예시:

pub struct FixedVec<T, W, E, B = Vec<W>> {
    bits: B,
    bit_width: usize,
    mask: W,
    len: usize,
}

이 구조는 signed 타입도, unsigned 타입도, 다양한 비트 폭과 메모리 타입(owned/borrowed)에 유연하게 적용할 수 있습니다.
또한 ZigZag 인코딩 같은 트릭으로 실수(음수/양수) 숫자도 똑똑하게 저장합니다.

마지막으로 Builder 패턴으로, 데이터의 최대값에 따라 자동으로 최적의 bit-width를 선택하거나, 원하는 값으로 직접 지정하는 것도 가능합니다.

정리: AI와 Rust에서 비트 패킹을 잘 쓰면 생기는 효용

엄청난 공간 절약: 데이터 전체에 쓸데없는 패딩 없이 필요한 만큼만 저장
빠른 무작위 접근: 소규모 비트폭일수록 캐시효과로 일반 벡터보다도 더 빠를 수 있음
유연한 확장성: 타입, 엔디안, 저장방식, 쓰기/읽기/반복 모두 원하는 대로 구현 가능
AI/인공지능 대형 데이터에도 딱: 작은 값의 집합, 특성상 데이터 압축이 필수적일 때 효과 극대화

비트 패킹은 대용량 AI 프로젝트와 Rust 개발에서 필수적인 무기! 실제 코드와 라이브러리는 아래 github에서 참고 가능하니, 직접 속도를 체험해보는 것도 추천합니다.

참고

[1] Engineering a fixed-width bit-packed integer vector in Rust - lukefleed.xyz

[2] SBVR: Summation of BitVector Representation for Efficient LLM Quantization - arXiv

이미지 출처

이미지 출처: Artem Podrez on Pexels