Làm chủ SIMD và Vectorization trong Rust: Tăng tốc phần cứng CPU hiện đại

Programming tutorial - IT technology blog
Programming tutorial - IT technology blog

Sự cố: Khi luồng nạp video 4K làm tê liệt cả cụm máy chủ

Mùa thu năm ngoái, đội ngũ chúng tôi chuyển đổi một dịch vụ xử lý hình ảnh sang Rust. Pipeline này chuẩn hóa các luồng video 4K liên tục (độ phân giải 3840×2160, tương đương khoảng 8,3 triệu pixel mỗi khung hình). Nó thực hiện căn chỉnh độ tương phản trên từng pixel và hòa trộn alpha cho bounding-box trên các mảng float 32-bit. Chúng tôi chọn Rust với kỳ vọng đạt được tốc độ native tương đương C mà không chịu bất kỳ chi phí runtime overhead nào.

Phiên bản triển khai ban đầu trông rất chỉn chu, đúng chuẩn idiomatic Rust:

pub fn normalize_pixel_luma(samples: &mut [f32], gain: f32) {
    for sample in samples.iter_mut() {
        *sample = (*sample * gain).clamp(-1.0, 1.0);
    }
}

Các bài benchmark nội bộ chạy rất mượt mà. Kiểm thử giả lập đều vượt qua trơn tru không chút lỗi lầm. Thế nhưng ngay khi lưu lượng production đạt 120 FPS trên 32 luồng camera RTSP, độ trễ xử lý mỗi khung hình vọt từ 3,2 ms lên 18,7 ms. Toàn bộ 16 vCPU bị nghẽn cứng ở mức 100%. Rayon đã phân phối các lát cắt dữ liệu đồng đều đến các worker thread, nhưng từng nhân CPU lại đang lãng phí chu kỳ máy để chạy các lệnh vô hướng (scalar). CPU chỉ đang xử lý từng float 32-bit một. Chúng tôi có thừa số nhân, nhưng lại bị thiếu hụt nghiêm trọng thông lượng xử lý trên từng nhân đơn lẻ.

Nguyên nhân gốc rễ: Vì sao vòng lặp vô hướng lại lãng phí tiềm năng phần cứng hiện đại

Các bộ vi xử lý máy chủ hiện đại đều tích hợp các đơn vị xử lý vector chuyên dụng song song với pipeline vô hướng. Một nhân x86_64 trang bị AVX2 cung cấp các thanh ghi vector 256-bit; AVX-512 mở rộng con số này lên tới 512-bit. ARM NEON trên các vi xử lý Graviton mang đến các luồng vector 128-bit. Thay vì chạy theo mô hình Dữ liệu đơn, Lệnh đơn (SISD), các đơn vị vector này thực thi theo mô hình Đa dữ liệu, Lệnh đơn (SIMD).

Một thanh ghi AVX2 256-bit đơn lẻ có thể chứa cùng lúc tám số thực float 32-bit (8 × 32 = 256 bits). Về lý thuyết, đơn vị thực thi có thể nhân cả tám giá trị này chỉ trong một chu kỳ xung nhịp. Tuy nhiên, file binary của chúng tôi lại đang chạy mã assembly vô hướng. Chúng tôi đã dịch ngược assembly đã biên dịch để kiểm tra chuyện gì đang xảy ra:

# Kiểm tra mã assembly đã biên dịch
cargo rustc --release -- --emit asm

Vòng lặp tải cao đầy rẫy các lệnh vô hướng vmulss thay vì các lệnh vector đóng gói vmulps. LLVM đã từ chối auto-vectorize vòng lặp này. Việc kết hợp phép nhân float với hàm clamp đã tạo ra các bước kiểm tra biên và các nhánh xử lý giá trị NaN tiềm ẩn. Do quy chuẩn IEEE-754 đòi hỏi việc lan truyền NaN phải được bảo toàn nghiêm ngặt, trình biên dịch không thể tự ý sắp xếp lại thứ tự tính toán mà không làm phá vỡ các cam kết tuân thủ chuẩn này. Các nhân AMD EPYC 7763 của chúng tôi bị bỏ phí tới 7/8 tiềm năng tính toán phần cứng.

So sánh các giải pháp trong môi trường Production

Suốt hai quý tiếp theo, chúng tôi đã đo đạc và đánh giá ba phương pháp vector hóa vòng lặp này trên khối lượng công việc thực tế.

1. Dựa vào cờ trình biên dịch và Auto-Vectorization

Lựa chọn đơn giản nhất là yêu cầu trình biên dịch tự động vector hóa thông qua các cờ tối ưu mạnh tay.

# .cargo/config.toml
[build]
rustflags = ["-C", "target-cpu=native", "-C", "llvm-args=-ffast-math"]
  • Ưu điểm: Không cần sửa đổi code. Giữ nguyên khả năng tương thích mã nguồn đa nền tảng.
  • Nhược điểm: Cực kỳ mong manh và dễ gãy. Một thay đổi nhỏ khi refactor, một phép kiểm tra biên thêm vào, hoặc logic đánh chỉ mục hơi khác biệt cũng có thể âm thầm phá vỡ thuật toán heuristic vector của LLVM. Trình biên dịch sẽ lặng lẽ quay về vòng lặp vô hướng mà không đưa ra bất kỳ cảnh báo nào.

2. Sử dụng Intrinsic phần cứng trực tiếp (std::arch::x86_64)

Rust cung cấp các intrinsic phần cứng trực tiếp qua core::arch. Điều này cho phép bạn viết thủ công các chỉ lệnh vector tương đương trực tiếp với mã assembly.

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

pub unsafe fn normalize_avx2(samples: &mut [f32], gain: f32) {
    let gain_vec = _mm256_set1_ps(gain);
    let min_vec = _mm256_set1_ps(-1.0);
    let max_vec = _mm256_set1_ps(1.0);
    
    let chunks = samples.chunks_exact_mut(8);
    let remainder = chunks.into_remainder();
    
    for chunk in chunks {
        let mut data = _mm256_loadu_ps(chunk.as_ptr());
        data = _mm256_mul_ps(data, gain_vec);
        data = _mm256_max_ps(min_vec, _mm256_min_ps(max_vec, data));
        _mm256_storeu_ps(chunk.as_mut_ptr(), data);
    }
    
    for item in remainder {
        *item = (*item * gain).clamp(-1.0, 1.0);
    }
}
  • Ưu điểm: Đạt thông lượng phần cứng tối đa với mã máy được sinh ra hoàn toàn có thể dự đoán được.
  • Nhược điểm: Bắt buộc phải dùng khối lệnh unsafe và phải bóc tách vòng lặp thủ công (loop peeling) cho các phần dư còn lại. Bạn cũng phải bao bọc lệnh gọi bằng tính năng kiểm tra CPU runtime (is_x86_feature_detected!("avx2")). Đoạn mã sẽ không thể biên dịch trên các phiên bản ARM64 Graviton nếu thiếu đi một bản cài đặt NEON riêng biệt thứ hai.

3. Vector hóa an toàn, đa nền tảng (Crate wide hoặc Nightly std::simd)

Portable SIMD cung cấp các kiểu vector không phụ thuộc vào phần cứng cụ thể như f32x8 và f32x4. Trên các máy chủ x86, mã nguồn sẽ tự biên dịch thành tập lệnh AVX2. Trên máy chủ ARM, nó tự động sinh ra tập lệnh NEON mà không cần tới bất kỳ khối unsafe nào.

use wide::f32x8;

pub fn normalize_portable_wide(samples: &mut [f32], gain: f32) {
    let gain_vec = f32x8::splat(gain);
    let min_vec = f32x8::splat(-1.0);
    let max_vec = f32x8::splat(1.0);
    
    let (chunks, remainder) = samples.as_chunks_mut::<8>();
    
    for chunk in chunks {
        let mut v = f32x8::from(*chunk);
        v = (v * gain_vec).fast_max(min_vec).fast_min(max_vec);
        *chunk = v.to_array();
    }
    
    for item in remainder {
        *item = (*item * gain).clamp(-1.0, 1.0);
    }
}
  • Ưu điểm: 100% Rust an toàn (safe Rust). Biên dịch mượt mà trên cả nền tảng x86 lẫn ARM với cú pháp rất súc tích.
  • Nhược điểm: Phải phụ thuộc vào thư viện ngoài (như wide hoặc pulp) trên bản Rust stable, hoặc phải dùng bộ toolchain nightly nếu muốn dùng std::simd.

Mô hình thực tế tại Production: SIMD đa luồng

Auto-vectorization hóa ra quá thiếu ổn định cho các pipeline CI/CD của chúng tôi, còn việc duy trì các triển khai intrinsic riêng biệt lại làm tăng gấp đôi gánh nặng kiểm thử. Chúng tôi quyết định chuẩn hóa bằng các mảng vector chia khối an toàn nhờ crate wide, kết hợp với Rayon để xử lý song song đa nhân tương tự như các mô hình concurrency khi chịu tải lớn.

Các bước triển khai cụ thể

Thêm các dependency sau vào file Cargo.toml:

[dependencies]
wide = "0.7"
rayon = "1.8"

Dưới đây là hàm xử lý thực tế trên production. Rayon chia nhỏ buffer để phân phối công việc đều qua các nhân, trong khi wide::f32x8 thực hiện vector hóa các thao tác ngay bên trong bộ nhớ đệm L1 cache của từng nhân:

use wide::f32x8;
use rayon::prelude::*;

pub fn parallel_simd_normalize(data: &mut [f32], factor: f32) {
    // 4096 số thực float = 16 KB, vừa vặn nằm gọn trong L1 data cache
    data.par_chunks_mut(4096).for_each(|slice| {
        let factor_vec = f32x8::splat(factor);
        let min_vec = f32x8::splat(-1.0);
        let max_vec = f32x8::splat(1.0);
        
        let chunks = slice.chunks_exact_mut(8);
        let remainder = chunks.into_remainder();
        
        for chunk in chunks {
            let mut vec = f32x8::new([
                chunk[0], chunk[1], chunk[2], chunk[3],
                chunk[4], chunk[5], chunk[6], chunk[7],
            ]);
            
            vec = (vec * factor_vec).fast_max(min_vec).fast_min(max_vec);
            
            let arr = vec.to_array();
            chunk.copy_from_slice(&arr);
        }
        
        // Xử lý các phần tử thừa còn lại không chia hết cho 8
        for item in remainder {
            *item = (*item * factor).clamp(-1.0, 1.0);
        }
    });
}

Kết quả đạt được trên Production sau 6 tháng

Giải pháp này đã vận hành ổn định trên production suốt sáu tháng qua. Dưới đây là những số liệu thực tế được ghi nhận từ Datadog:

  • Độ trễ: Thời gian xử lý một khung hình 4K giảm ngoạn mục từ 18,7 ms xuống chỉ còn 2,4 ms—thông lượng tăng gấp 7,8 lần.
  • Mức độ sử dụng CPU: Mức tải CPU trung bình của cả cụm giảm sâu từ 100% xuống còn 32%. Chúng tôi đã cắt giảm số lượng instance EC2 từ hai mươi máy c5.4xlarge xuống còn mười máy, tiết kiệm khoảng 2.400 USD mỗi tháng.
  • Khả năng chuyển đổi đa kiến trúc: Khi chuyển các worker node sang AWS Graviton3 (c7g.4xlarge, ARM NEON), cùng một đoạn mã nguồn đã biên dịch và vận hành trơn tru mà không cần chỉnh sửa bất kỳ dòng nào, mở ra tiềm năng tích hợp cho cả các hệ thống mở rộng như WebAssembly với Rust. Hiệu năng đạt được vẫn duy trì trong khoảng chênh lệch 5% so với kết quả benchmark AVX-512 trên x86.

Trình biên dịch rất thông minh, nhưng chúng không thể tự suy đoán được ý đồ kiến trúc của bạn qua những quy tắc khắt khe của chuẩn float IEEE. Việc chủ động vector hóa tường minh sẽ loại bỏ phỏng đoán của trình biên dịch, biến các đường ống tính toán phần cứng đang nhàn rỗi thành thông lượng thực tế, và giữ cho các pipeline tải nặng luôn vận hành bền bỉ với hiệu năng ổn định tuyệt đối.

Share: