Rust의 SIMD 추상화가 내부의 unsafe를 크게 줄이면서도 안전하고 사용하기 쉬운 API를 제공하는 방법을 살펴봅니다.
9분 읽기
2일 전
Rust의SIMD 추상화는 제가 바라는 만큼 안전하지 않았습니다. 지금까지는요.
로우 SIMD intrinsic은 쓰기 불편하다는 건 비밀도 아닙니다.
우리가 쓰고 싶은 것은 a + b이지, 이런 괴물 같은 것이 아닙니다:
unsafe {
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "avx2"))]
_mm256_add_ps(a, b)
#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "sse", not(target_feature = "avx2")))]
_mm_add_ps(a, b)
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
vaddq_f32(a, b)
}
이걸 보세요. 끔찍합니다. 게다가 전체가 unsafe로 감싸여 있습니다!
그리고 이건 단순화한 예시일 뿐입니다. 여전히 다음은 처리하지 못합니다:
&[f32] 같은 데이터를 각 intrinsic이 받는 형태로 실제로 로드하기다행히 Rust는 이 모든 것을 처리해 주고, 그냥 a + b라고 쓸 수 있게 해 주는 다양한 SIMD 추상화를 제공합니다.
단 한 가지 문제가 있습니다. 내부는 여전히 unsafe로 가득합니다. unsafe가 사라진 것이 아니라 숨겨졌을 뿐입니다. 표면 바로 아래에 엄청난 양이 숨어 있고, 가끔씩 망가지곤하고있었습니다.
정확히 말하면, 그랬습니다. 지금까지는요.
unsafe가 필요할까?아주 오랫동안 _mm256_add_ps 같은 각 intrinsic 함수를 호출할 때 unsafe로 감싸는 것을 피할 수 없었습니다. 현재 실행 중인 CPU에서 사용할 수 없을 때 그런 함수를 호출하는 것은 불법이기 때문입니다.
그래서 각 intrinsic에 어떤 명령어가 필요한지, 내가 어떤 명령어를 사용할 수 있는지 추적하고, 둘을 대조해서 주어진 함수를 안전하게 호출할 수 있는지 판단하는 메커니즘이 반드시 필요했습니다.
이 작업은 손으로 하면 지루하고, 코드 생성기로 하면 복잡했으며, 언제나 실수하기 쉬웠고, 모든 intrinsic 주위에 unsafe가 필요했습니다.
이 상황은 Rust 1.87에서 바뀌었습니다. 컴파일러가 필요한 명령어 집합을 직접 추적하기 시작해서, 이제 이렇게 쓸 수 있게 되었기 때문입니다:
#[target_feature(enable = "avx2")]
fn add_avx2(a: __m256, b: __m256) -> __m256 {
_mm256_add_ps(a, b)
}
보세요, unsafe가 없습니다!
…아직은요.
그래도 이것만으로는 a + b를 쓸 수 없습니다. 할 수 있는 최선은 이것입니다:
unsafe { add_avx2(a, b) }
이건 단지 unsafe를 한 단계 위로 올려놓은 것뿐입니다. 이제는 올바른 #[target_feature]가 붙은 함수 내부에서 intrinsic을 호출할 수 있지만, 호출 사슬 어딘가에는 여전히 unsafe가 있어야 합니다.
또 다른 문제는 더 근본적입니다. 여러분 타입에 대한 + 구현에 #[target_feature]를 붙일 수 없습니다. +는 언제나 사용 가능해야 하기 때문입니다. 그러니 이 메커니즘만으로는 a + b를 쓸 수 없습니다.
최종 해법이 어떻게 동작하는지 이해하려면, 먼저 CPU 기능 감지가 어떻게 작동하는지 이해해야 합니다.
보통 AVX2 같은 CPU 기능 확인은 is_x86_feature_detected!("avx2")를 사용해 런타임에 수행합니다. 하지만 두 수를 더할 때마다 이 검사를 매번 실행하고 싶지는 않습니다. 그러면 성능이 완전히 망가집니다. 우리는 이 검사를 한 번만 하고, 그 시점부터 AVX2 명령어를 사용해도 안전하다는 것을 컴파일러에 증명하고 싶습니다.
대신 이 증명을 타입 시스템에 _위조 불가능한 토큰_으로 인코딩할 수 있습니다. 이는 내부 필드가 private인 크기 0 타입입니다. 이 토큰을 얻는 유일한 방법은 CPU 기능 검사를 수행하는 함수를 호출하는 것입니다. 검사가 통과하면 함수가 토큰을 건네줍니다:
pub struct Avx2(()); fn detect_avx2() -> Option<Avx2> {
if is_x86_feature_detected!("avx2") {
Some(Avx2(()))
} else {
None
}
}
그리고 이것은 크기 0 타입이므로, 이 토큰을 여기저기 전달해도 런타임 오버헤드는 없습니다. 순수하게 컴파일 타임 증명으로만 존재합니다.
핵심은 Avx2 구조체의 인스턴스를 가지고 있는 한, 시스템에서 AVX2 명령어를 사용할 수 있다고 확신할 수 있다는 점입니다.
컴파일러는 모르지만, 이 함수는 안전하게 호출할 수 있습니다:
#[target_feature(enable = "avx2")]
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
_mm256_add_ps(a, b)
}
이 함수는 Avx2 토큰이 있을 때만 호출할 수 있고, 그 토큰은 시스템에서 AVX2 명령어를 사용할 수 있을 때만 얻을 수 있습니다.
이것이 타당하다는 점을 컴파일러에게 unsafe를 사용해 설명할 수 있다면, 그 unsafe를 단 한 번만 작성하고 어디서나 재사용할 수 있습니다.
우리에게 필요한 것은 안전하게 호출할 수 있는 매크로입니다:
with_avx2!(
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
_mm256_add_ps(a, b)
}
) 하지만 내부적으로는 다음과 같이 확장됩니다:
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
unsafe { inner(token, a, b) } #[target_feature(enable = "avx2")]
fn inner(token: Avx2, a: __m256, b: __m256) -> __m256 {
_mm256_add_ps(a, b)
}
}
이제 AVX2에 없는 intrinsic을 사용하면 컴파일러가 거부합니다!
맞춤형 target feature 추적 없이도 SIMD intrinsic에 대한 안전한 프로그래밍 인터페이스를 제공하는 데 성공한 것입니다!
내부에 unsafe 블록 하나가 여전히 있긴 하지만, 이것은 sound한 API 안에 캡슐화되어 있으므로 이를 오용해 메모리 안전성 버그를 일으킬 수 없습니다. 그런 의미에서 이것은 println!과 마찬가지로, unsafe 코드를 안전하게 추상화합니다.
이 작성자의 업데이트를 받으려면 Medium에 무료로 가입하세요.
더 빠른 로그인을 위해 나를 기억하기
이 방식을 사용하면 검토하고 감사해야 하는 대상은 _이 매크로 하나뿐_이며, 제각각 만들어진 수백 수천 개의 unsafe 블록이 아닙니다. 그리고 구현에서 우리가 실제로 망칠 수 있는 것은 오직 다음뿐입니다:
#[target_feature]에 매핑하기unsafe fn이 안전한 컨텍스트에서 호출되도록 허용하기그리고 이 두 실패 모드는 모두 확인하기 꽤 쉽습니다.
이제 unsafe 없이 add_avx2(token, a, b)를 호출할 수 있게 되었지만, 그래도 아직 a + b에는 도달하지 못했습니다. 그건 어떻게 해결할까요?
a + b의 구현에 #[target_feature]를 붙일 수는 없습니다. 어디에서나 안전하게 호출할 수 있어야 하기 때문입니다. 그리고 함수가 a와 b는 받지만 token은 받지 않으므로, 토큰을 함수에 전달할 수도 없습니다.
하지만 설령 그렇게 할 수 있다 해도, API가 꽤 보기 흉해집니다. 우리는 사용자가 토큰을 신경 쓰지 않아도 a + b가 항상 동작하고 자동으로 최적의 SIMD 명령어를 사용하길 원합니다.
이 두 문제를 한 번에 해결하기 위해 제네릭을 사용할 수 있습니다. 사용 가능한 명령어 집합에 대해 제네릭한 f32x8 타입을 정의하면, 그 위에 덧셈을 구현하면서 내부에 토큰을 몰래 넣을 수도 있고, 각 SIMD 명령어 집합마다 별도의 구현도 만들 수 있습니다!
모양은 이렇습니다:
pub trait Level {} #[derive(Clone, Copy)]
pub struct Avx2(());
impl Level for Avx2 {}
pub struct f32x8<L: Level> {
data: [f32; 8],
token: L,
}
impl std::ops::Add for f32x8<Avx2> {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
let result = add_avx2(self.token, self, rhs);
Self {
data: store_m256(result),
token: self.token,
}
}
}
그리고 다른 어떤 명령어 집합에 대해서도, 또는 SIMD를 전혀 사용할 수 없을 때도 똑같이 쉽게 동작하게 만들 수 있습니다:
#[derive(Clone, Copy)]
pub struct NoSimd(());
impl Level for NoSimd {} impl std::ops::Add for f32x8<NoSimd> {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
let result = std::array::from_fn(|i| self.data[i] + rhs.data[i]);
Self {
data: result,
token: self.token,
}
}
}
이렇게 해서 안전성과 런타임 명령어 선택을 한 번에 해결했습니다!
시스템에서 사용 가능한 최선의 Level을 제공하는 편의 함수를 하나 추가하면, SIMD에 대해 거의 완벽한 API를 얻게 됩니다!
불행히도 a + b를 작성하고 그것이 SIMD 명령어로 낮아지게 하는 데에는 근본적인 문제가 있습니다. 바로 함수 호출 오버헤드입니다.
함수 호출은 공짜는 아니지만 꽤 저렴합니다. 겨우 몇 개의 CPU 명령어면 됩니다. 하지만 _몇 개_의 명령어도, 우리가 덧셈 구현에 방금 사용한 _한 개_의 명령어에 비하면 훨씬 많습니다!
그래서 중간에 함수 호출이 끼어 있으면 덧셈 성능은 급락합니다. 그리고 성능이야말로 애초에 SIMD를 쓰는 이유 전부입니다!
보통 컴파일러는 인라이닝을 통해 이 오버헤드를 지우는 데 꽤 능숙합니다. 호출한 함수의 구현을 그 함수를 호출하는 함수 안으로 사실상 복사해 넣어서, 더 이상 함수도 없고 오버헤드도 없게 만듭니다.
하지만 #[target_feature] 애너테이션은 여기에 훼방을 놓습니다. 컴파일러는 #[target_feature] 애너테이션이 붙은 함수를, 그것이 붙어 있지 않은 함수 안으로 인라인할 수 없습니다. 그 함수 안에서는 필요한 기능을 사용할 수 없기 때문입니다!
그리고 #[target_feature] 애너테이션을 붙일 수 없는 것이 무엇이냐고요? 네, 맞습니다.
그렇다면 도대체 어떻게 a + b가 SIMD와 함께 동작하게 만들 수 있을까요?
a + b를 구현하는 함수에 #[target_feature]를 붙일 수는 없지만, a + b를 호출하는 함수에는 붙일 수 있습니다!
그 다음 인라이닝을 사용해 a + b의 구현이 그것을 호출하는 함수 안으로 복사되도록 만들면, 최종적으로 #[target_feature] 컨텍스트 안에 들어가게 됩니다.
그러면 호출 사슬은 이렇게 됩니다:
#[target_feature(enable = "avx2")]
fn do_stuff() {
c = a + b;
} #[inline(always)]
fn add(self, rhs: Self) -> Self::Output {
add_avx2(self.token, self, rhs);
}
#[inline]
#[target_feature(enable = "avx2")]
fn add_avx2(token: Avx2, a: __m256, b: __m256) -> __m256 {
_mm256_add_ps(a, b)
}
이 방식은 동작합니다.
올바른 위치에 이런 애너테이션을 추가하고 SIMD 레벨에 대해 추상화하면, 매크로조차 필요하지 않습니다.
문제는 SIMD 타입에 대해 a + b를 호출할 때마다, 그것을 #[inline(always)] 또는 #[target_feature]가 붙은 함수에서 호출해야 한다는 점입니다. 그렇지 않으면 코드는 여전히 컴파일되지만 성능이 급락합니다.
직접 보고 싶으신가요? 이 예제를 열고 #[target_feature]를 제거한 다음, 생성되는 어셈블리가 완전한 공포 쇼로 바뀌는 것을 보세요.
이 문제를 어떻게 해결할 수 있을지는 잘 모르겠습니다. SIMD로 a + b를 구현하는 어떤 접근법에 대해서도 이 제한은 꽤 근본적인 것처럼 보입니다.
Struct Target Features RFC는 add_avx2(token, a, b)와 add<S: Simd>(token, a, b)에 대해서는 이것을 해결하지만, a + b로 가는 길은 아직 보이지 않습니다.
모든 SIMD 코드에 내재한 인라이닝 문제에도 불구하고, 우리는 놀라울 정도로 적은, 전례 없이 낮은 수준의 unsafe 코드만으로도 상당히 쾌적한 SIMD 추상화를 제공하는 데 성공했습니다.
이 아이디어의 프로덕션 버전은 fearless_simd v0.5에서 찾을 수 있습니다. 지금 바로 사용 가능하며, 가까운 패키지 레지스트리에서 만나볼 수 있습니다! 그리고 프로덕션에서 이것들이 어떻게 함께 맞물리는지 보려면 이 작은 예제를 확인해 보세요.
이를 구현하는 데 사용된 매크로도 공개되어 있으므로, 하드웨어를 최대한 활용하기 위해 a + b 같은 고수준 연산과 플랫폼별 intrinsic을 쉽게 섞어 쓸 수 있습니다.
fearless_simd에는 unsafe 블록이 하나보다 더 있습니다. safe_unaligned_simd crate의 기능도 제공하기 때문입니다. 하지만 그 부분 역시 원래 것보다 훨씬 더 적은 양의 unsafe 코드로 구현되었습니다.
저에게 고수준 SIMD 추상화를 사용하는 데 대한 장벽은 언제나 그것들이 끌고 오는 엄청난 양의 unsafe였습니다. 무서웠고, 정당화하기 어려웠습니다.
하지만 이제 Rust의 SIMD는 진정으로 fearless할 수 있습니다.
제가 알기로는 제가 이것을 처음으로 프로덕션에 넣은 사람이라는 사실이 놀랍습니다. 왜냐하면 이 아이디어를 떠올린 첫 번째 사람이 분명 제가 아니기 때문입니다.
CPU 기능 토큰은 오래되고 흔한 아이디어입니다. pulp crate는 이것을 수년간 사용해 왔지만, intrinsic 주위에 손으로 작성한 unsafe 래퍼에 의존했고, 가끔은 그것을 잘못 작성하기도 했습니다.
제네릭을 사용해 여러 구현을 생성하는 것 역시 오래된 아이디어입니다. 이것은 8년 전 원래 fearless_simd 개념의 일부였습니다. 그것보다 더 오래된 simdeez crate도 비슷한 무언가를 사용하는 것처럼 보입니다.
토큰과 rustc에 위임하는 단일 안전 래퍼를 결합하는 핵심 통찰도 저만의 것은 아닙니다. fearless_simd crate의 맥락만 보더라도, Raph Levien이 이를 실험해 본 적이 있고, Daniel McNab은 몇 달 전에 제 것보다 더 정교한 구현을 만들었습니다.
Daniel의 접근법은 fearless_simd가 사용하는 소수의 고정된 CPU 기능 레벨과 달리, 개별 CPU 기능 하나하나를 세밀하게 추적할 수 있게 해 줍니다. 표현력은 더 높지만 복잡성이라는 대가가 있었고, 다른 메인테이너가 리뷰를 맡지 않아서 그의 접근법은 결국 병합되지 못했습니다. 언젠가 이것이 독립적인 crate로 공개되기를 여전히 바랍니다.
Daniel과, fearless_simd에 대한 제 모든 PR을 리뷰해 준 Laurenz Stampfl에게 감사드립니다. PR들이 컸는데도 빠르게 리뷰해 준 점을 정말 고맙게 생각합니다!