Skip to content

Commit 63366b3

Browse files
5000user50005000user5000
committed
perf: add parallel SIMD implementation with blocking for both float and LUT versions
1 parent 3926b26 commit 63366b3

2 files changed

Lines changed: 179 additions & 71 deletions

File tree

‎Makefile‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,12 @@ TEST_SRC := $(TEST_DIR)/run_benchmark.cpp
2424
TARGET_MAIN := $(BUILD_DIR)/run_benchmark
2525

2626
# correctness suite
27-
CORR_SRC := $(TEST_DIR)/test_correctness.cpp # fix filename
28-
TARGET_CORR := $(BUILD_DIR)/test_correctness # fix binary name
27+
CORR_SRC := $(TEST_DIR)/test_correctness.cpp
28+
TARGET_CORR := $(BUILD_DIR)/test_correctness
29+
30+
# matrix ops test
31+
MATRIX_OPS_SRC := $(TEST_DIR)/test_matrix_ops.cpp
32+
TARGET_MATRIX_OPS := $(BUILD_DIR)/test_matrix_ops
2933

3034
HEADERS := \
3135
$(SRC_DIR)/layout_policies.hpp \
@@ -35,24 +39,28 @@ HEADERS := \
3539
$(SRC_DIR)/lut_utils.hpp \
3640
$(SRC_DIR)/post_processing.hpp
3741

38-
.PHONY: all run test clean pytest
42+
.PHONY: all run test clean pytest matrix_ops matrix_ops_float matrix_ops_lut
3943

40-
all: $(BUILD_DIR) $(TARGET_MAIN) $(TARGET_CORR) mpgemm$(PYEXT)
44+
all: $(BUILD_DIR) $(TARGET_MAIN) $(TARGET_CORR) $(TARGET_MATRIX_OPS) mpgemm$(PYEXT)
4145

4246
# ensure build directory exists
4347
$(BUILD_DIR):
4448
mkdir -p $(BUILD_DIR)
4549

4650
# build main
47-
$(TARGET_MAIN): $(TEST_SRC) $(HEADERS)
51+
$(TARGET_MAIN): $(BUILD_DIR) $(TEST_SRC) $(HEADERS)
4852
$(CXX) $(CXXFLAGS) $(TEST_SRC) -o $(TARGET_MAIN) $(LDFLAGS) $(LDLIBS)
4953

5054
# build correctness suite
51-
$(TARGET_CORR): $(CORR_SRC) $(HEADERS)
55+
$(TARGET_CORR): $(BUILD_DIR) $(CORR_SRC) $(HEADERS)
5256
$(CXX) $(CXXFLAGS) $(CORR_SRC) -o $(TARGET_CORR) $(LDFLAGS) $(LDLIBS)
5357

58+
# build matrix ops test
59+
$(TARGET_MATRIX_OPS): $(BUILD_DIR) $(MATRIX_OPS_SRC) $(HEADERS)
60+
$(CXX) $(CXXFLAGS) -pthread $(MATRIX_OPS_SRC) -o $(TARGET_MATRIX_OPS) $(LDFLAGS) $(LDLIBS)
61+
5462
# build pybind11 module
55-
mpgemm$(PYEXT): src/bindings.cpp $(HEADERS)
63+
mpgemm$(PYEXT): $(BUILD_DIR) src/bindings.cpp $(HEADERS)
5664
$(CXX) $(CXXFLAGS) $(PYBIND11_INC) -fPIC -shared src/bindings.cpp -o $@ $(LDFLAGS) $(LDLIBS)
5765

5866
# run pytest
@@ -65,6 +73,19 @@ run: all
6573
test: $(TARGET_CORR)
6674
./$(TARGET_CORR)
6775

76+
matrix_ops: $(TARGET_MATRIX_OPS)
77+
@echo "Running float version..."
78+
@./$(TARGET_MATRIX_OPS) float
79+
@echo "\nRunning LUT version..."
80+
@./$(TARGET_MATRIX_OPS) lut
81+
82+
# 方便的命令
83+
matrix_ops_float: $(TARGET_MATRIX_OPS)
84+
./$(TARGET_MATRIX_OPS) float
85+
86+
matrix_ops_lut: $(TARGET_MATRIX_OPS)
87+
./$(TARGET_MATRIX_OPS) lut
88+
6889
clean:
6990
rm -rf $(BUILD_DIR)
7091
rm -f mpgemm$(PYEXT)

‎src/matrix_ops.hpp‎

Lines changed: 151 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
#include <type_traits>
77
#include <vector>
88
#include <immintrin.h>
9+
#include <thread>
10+
#include <mutex>
911

1012
// =============================================================
1113
// Helper: unpack a Matrix<> that uses Int4Storage into a
@@ -29,82 +31,57 @@ std::vector<uint8_t> unpack_int4(const Mat4& M)
2931
}
3032

3133
// =============================================================
32-
// Naive reference GEMM (kept unchanged for correctness checks)
34+
// High-performance parallel GEMM implementation
35+
// Supports any numeric type through templates
3336
// =============================================================
3437

3538
template<typename MA, typename MB>
36-
auto matmul(const MA& A, const MB& B)
39+
auto matmul(const MA& A, const MB& B, size_t num_threads = 4)
3740
{
3841
using T = decltype(A.at(0, 0));
3942
static_assert(std::is_same_v<T, decltype(B.at(0, 0))>, "Element types must match");
4043

4144
size_t M = A.rows(), K = A.cols(), N = B.cols();
4245
Matrix<T, RowMajor, PlainStorage<T>> C(M, N);
46+
std::vector<std::thread> threads;
47+
48+
// 計算每個執行緒處理的行數
49+
size_t rows_per_thread = (M + num_threads - 1) / num_threads;
50+
51+
// 為每個執行緒分配工作
52+
for (size_t t = 0; t < num_threads; ++t) {
53+
threads.emplace_back([&, t]() {
54+
size_t start_row = t * rows_per_thread;
55+
size_t end_row = std::min(start_row + rows_per_thread, M);
56+
57+
// 為每個執行緒創建局部結果矩陣
58+
Matrix<T, RowMajor, PlainStorage<T>> local_C(M, N);
59+
60+
for (size_t i = start_row; i < end_row; ++i) {
61+
for (size_t k = 0; k < K; ++k) {
62+
T a = A.at(i, k);
63+
for (size_t j = 0; j < N; ++j) {
64+
local_C.set(i, j, local_C.at(i, j) + a * B.at(k, j));
65+
}
66+
}
67+
}
4368

44-
for (size_t i = 0; i < M; ++i)
45-
for (size_t k = 0; k < K; ++k) {
46-
T a = A.at(i, k);
47-
for (size_t j = 0; j < N; ++j)
48-
C.set(i, j, C.at(i, j) + a * B.at(k, j));
49-
}
50-
return C;
51-
}
52-
53-
// =============================================================
54-
// High‑speed LUT GEMM — expects *unpacked* uint8 buffers.
55-
// * Au shape: M × K contiguous
56-
// * Bu shape: K × N contiguous
57-
// Works with or without AVX2 (scalar fallback).
58-
// =============================================================
59-
60-
auto matmul_lut_fast(const std::vector<uint8_t>& Au,
61-
const std::vector<uint8_t>& Bu,
62-
size_t M, size_t K, size_t N,
63-
const ProductLookupTable<uint8_t, uint8_t, int32_t>& lut)
64-
{
65-
const int32_t* lut_ptr = lut.data();
66-
const int32_t stride = static_cast<int32_t>(lut.row_stride());
67-
68-
#if defined(__AVX2__)
69-
const __m256i vstride = _mm256_set1_epi32(stride);
70-
#endif
71-
72-
Matrix<int32_t, RowMajor, PlainStorage<int32_t>> C(M, N);
73-
74-
for (size_t i = 0; i < M; ++i) {
75-
const uint8_t* rowA = &Au[i * K];
76-
for (size_t j = 0; j < N; ++j) {
77-
int32_t acc = 0;
78-
size_t k = 0;
79-
80-
#if defined(__AVX2__)
81-
// --- AVX2: process 8 elements (k dimension) per iteration ---
82-
for (; k + 7 < K; k += 8) {
83-
// load 8 uint8 from A row / B column (row‑major & col‑major buffers)
84-
__m128i w8 = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(rowA + k));
85-
__m128i a8 = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(&Bu[k * N + j]));
86-
87-
__m256i w32 = _mm256_cvtepu8_epi32(w8);
88-
__m256i a32 = _mm256_cvtepu8_epi32(a8);
89-
__m256i idx = _mm256_add_epi32(_mm256_mullo_epi32(w32, vstride), a32);
90-
__m256i vals = _mm256_i32gather_epi32(lut_ptr, idx, 4);
91-
92-
// horizontal sum of 8 lanes
93-
__m128i low = _mm256_castsi256_si128(vals);
94-
__m128i high = _mm256_extracti128_si256(vals, 1);
95-
__m128i sum = _mm_add_epi32(low, high);
96-
sum = _mm_hadd_epi32(sum, sum);
97-
sum = _mm_hadd_epi32(sum, sum);
98-
acc += _mm_cvtsi128_si32(sum);
69+
// 合併結果
70+
static std::mutex mtx;
71+
std::lock_guard<std::mutex> lock(mtx);
72+
for (size_t i = start_row; i < end_row; ++i) {
73+
for (size_t j = 0; j < N; ++j) {
74+
C.set(i, j, C.at(i, j) + local_C.at(i, j));
75+
}
9976
}
100-
#endif
101-
// --- scalar remainder (or full loop if no AVX2) ---
102-
for (; k < K; ++k)
103-
acc += lut_ptr[rowA[k] * stride + Bu[k * N + j]];
77+
});
78+
}
10479

105-
C.set(i, j, acc);
106-
}
80+
// 等待所有執行緒完成
81+
for (auto& thread : threads) {
82+
thread.join();
10783
}
84+
10885
return C;
10986
}
11087

@@ -149,4 +126,114 @@ Matrix<T> matmul_mkl(const Matrix<T>& A, const Matrix<T>& B) {
149126
}
150127
return C;
151128
}
152-
#endif
129+
#endif
130+
131+
// =============================================================
132+
// High‑speed LUT GEMM — expects *unpacked* uint8 buffers.
133+
// * Au shape: M × K contiguous
134+
// * Bu shape: K × N contiguous
135+
// Works with or without AVX2 (scalar fallback).
136+
// =============================================================
137+
138+
auto matmul_lut_fast(const std::vector<uint8_t>& Au,
139+
const std::vector<uint8_t>& Bu,
140+
size_t M, size_t K, size_t N,
141+
const ProductLookupTable<uint8_t, uint8_t, int32_t>& lut,
142+
size_t block_size = 64, // 增加區塊大小以減少同步開銷
143+
size_t num_threads = 4)
144+
{
145+
const int32_t* lut_ptr = lut.data();
146+
const int32_t stride = static_cast<int32_t>(lut.row_stride());
147+
Matrix<int32_t, RowMajor, PlainStorage<int32_t>> C(M, N);
148+
std::vector<std::thread> threads;
149+
std::mutex mtx;
150+
151+
// 預取 LUT 到 L1 cache
152+
for (size_t i = 0; i < 16; ++i) {
153+
for (size_t j = 0; j < 16; ++j) {
154+
_mm_prefetch(reinterpret_cast<const char*>(&lut_ptr[i * stride + j]), _MM_HINT_T0);
155+
}
156+
}
157+
158+
// 計算每個執行緒處理的行數
159+
size_t rows_per_thread = (M + num_threads - 1) / num_threads;
160+
161+
// 為每個執行緒分配工作
162+
for (size_t t = 0; t < num_threads; ++t) {
163+
threads.emplace_back([&, t]() {
164+
size_t start_row = t * rows_per_thread;
165+
size_t end_row = std::min(start_row + rows_per_thread, M);
166+
167+
// 為每個執行緒創建局部結果矩陣
168+
Matrix<int32_t, RowMajor, PlainStorage<int32_t>> local_C(M, N);
169+
170+
// 分塊處理
171+
for (size_t i = start_row; i < end_row; i += block_size) {
172+
size_t i_end = std::min(i + block_size, end_row);
173+
174+
for (size_t j = 0; j < N; j += block_size) {
175+
size_t j_end = std::min(j + block_size, N);
176+
177+
for (size_t k = 0; k < K; k += block_size) {
178+
size_t k_end = std::min(k + block_size, K);
179+
180+
// 處理當前區塊
181+
for (size_t ii = i; ii < i_end; ++ii) {
182+
const uint8_t* rowA = &Au[ii * K];
183+
184+
for (size_t jj = j; jj < j_end; ++jj) {
185+
int32_t acc = 0;
186+
187+
#if defined(__AVX2__)
188+
// 使用 AVX2 處理 8 個元素
189+
for (size_t kk = k; kk + 7 < k_end; kk += 8) {
190+
__m128i w8 = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(rowA + kk));
191+
__m128i a8 = _mm_loadl_epi64(reinterpret_cast<const __m128i*>(&Bu[kk * N + jj]));
192+
193+
__m256i w32 = _mm256_cvtepu8_epi32(w8);
194+
__m256i a32 = _mm256_cvtepu8_epi32(a8);
195+
__m256i idx = _mm256_add_epi32(_mm256_mullo_epi32(w32, _mm256_set1_epi32(stride)), a32);
196+
197+
// 預取下一個 LUT 值
198+
_mm_prefetch(reinterpret_cast<const char*>(&lut_ptr[_mm256_extract_epi32(idx, 0)]), _MM_HINT_T0);
199+
200+
__m256i vals = _mm256_i32gather_epi32(lut_ptr, idx, 4);
201+
202+
// 水平加總
203+
__m128i low = _mm256_castsi256_si128(vals);
204+
__m128i high = _mm256_extracti128_si256(vals, 1);
205+
__m128i sum = _mm_add_epi32(low, high);
206+
sum = _mm_hadd_epi32(sum, sum);
207+
sum = _mm_hadd_epi32(sum, sum);
208+
acc += _mm_cvtsi128_si32(sum);
209+
}
210+
#endif
211+
// 處理剩餘元素
212+
for (size_t kk = k + ((k_end - k) & ~7); kk < k_end; ++kk) {
213+
acc += lut_ptr[rowA[kk] * stride + Bu[kk * N + jj]];
214+
}
215+
216+
local_C.set(ii, jj, local_C.at(ii, jj) + acc);
217+
}
218+
}
219+
}
220+
}
221+
}
222+
223+
// 合併結果
224+
std::lock_guard<std::mutex> lock(mtx);
225+
for (size_t i = start_row; i < end_row; ++i) {
226+
for (size_t j = 0; j < N; ++j) {
227+
C.set(i, j, C.at(i, j) + local_C.at(i, j));
228+
}
229+
}
230+
});
231+
}
232+
233+
// 等待所有執行緒完成
234+
for (auto& thread : threads) {
235+
thread.join();
236+
}
237+
238+
return C;
239+
}

0 commit comments

Comments
 (0)