Skip to content

Commit 3b4fca1

Browse files
authored
ggml-cpu: Enable tiled matmul on AIX (ggml-org#25199)
The matmul_tiled path uses large local stack buffers for A_pack and B_pack. On AIX this can trigger a segmentation fault, so reduce the buffer footprint there to keep the tiled path usable. Performance Impact: ~ 2x gains in PP_Speed for FP32, Q4_0 and Q8_0 models tested with llama-bench, llama-batched-bench and llama-cli. Models used: Llama3.2 3b Instruct F32, qwen 2.5 3b Q4_0 and Q8_0
1 parent 86961ef commit 3b4fca1

1 file changed

Lines changed: 20 additions & 14 deletions

File tree

ggml/src/ggml-cpu/llamafile/sgemm.cpp

Lines changed: 20 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2321,24 +2321,28 @@ class tinyBLAS_Q0_PPC {
23212321
}
23222322

23232323
void matmul(int64_t m, int64_t n) {
2324-
#if defined(_AIX) || defined(__BIG_ENDIAN__)
2325-
mnpack(0, m, 0, n);
2326-
#else
2327-
const int64_t mc = 64;
2328-
const int64_t kc = 64;
2324+
int64_t mc = 64;
23292325
int64_t nc = 64;
2326+
int64_t kc = 64;
2327+
int64_t n_chunk = 64;
2328+
#if defined(_AIX) || defined(__BIG_ENDIAN__)
2329+
mc = 32;
2330+
nc = 32;
2331+
kc = 32;
2332+
n_chunk = 32
2333+
#endif
23302334
int64_t n_aligned = 0;
2331-
if (n % 64 == 0) {
2335+
if (n % n_chunk == 0) {
23322336
n_aligned = n;
23332337
} else if (n == 4) {
23342338
n_aligned = 4;
2335-
} else if (n < 64) {
2339+
} else if (n < n_chunk) {
23362340
n_aligned = (n / 8) * 8;
23372341
} else {
2338-
n_aligned = (n / 64) * 64;
2342+
n_aligned = (n / n_chunk) * n_chunk;
23392343
}
23402344
if (n_aligned > 0) {
2341-
if (n_aligned % 64 == 0) nc = 64;
2345+
if (n_aligned % n_chunk == 0) nc = n_chunk;
23422346
else if (n_aligned == n) nc = n;
23432347
else if (n_aligned % 32 == 0) nc = 32;
23442348
else if (n_aligned % 24 == 0) nc = 24;
@@ -2354,7 +2358,6 @@ class tinyBLAS_Q0_PPC {
23542358
} else {
23552359
mnpack(0, m, 0, n);
23562360
}
2357-
#endif
23582361
}
23592362

23602363
private:
@@ -3195,16 +3198,19 @@ class tinyBLAS_PPC {
31953198
}
31963199

31973200
void matmul(int64_t m, int64_t n) {
3201+
int64_t mc = 256;
3202+
int64_t nc = 256;
3203+
int64_t kc = 256;
31983204
#if defined(_AIX) || defined(__BIG_ENDIAN__)
3199-
mnpack(0, m, 0, n);
3200-
#else
3201-
int64_t mc = 256; int64_t nc = 256; int64_t kc = 256;
3205+
mc = 128;
3206+
nc = 128;
3207+
kc = 128;
3208+
#endif
32023209
if (m % mc == 0 && n % nc == 0 && k % kc == 0) {
32033210
matmul_tiled(m, n, mc, nc, kc);
32043211
} else {
32053212
mnpack(0, m, 0, n);
32063213
}
3207-
#endif
32083214
}
32093215

32103216
private:

0 commit comments

Comments
 (0)