diff --git a/ggml/src/ggml-rocket/CMakeLists.txt b/ggml/src/ggml-rocket/CMakeLists.txt index 19ec91a2e..71fe5bca0 100644 --- a/ggml/src/ggml-rocket/CMakeLists.txt +++ b/ggml/src/ggml-rocket/CMakeLists.txt @@ -19,3 +19,6 @@ target_include_directories(ggml-rocket PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ) + +find_package(Threads REQUIRED) +target_link_libraries(ggml-rocket PRIVATE Threads::Threads) diff --git a/ggml/src/ggml-rocket/ggml-rocket.cpp b/ggml/src/ggml-rocket/ggml-rocket.cpp index 6c18aa97a..20e57f8e3 100644 --- a/ggml/src/ggml-rocket/ggml-rocket.cpp +++ b/ggml/src/ggml-rocket/ggml-rocket.cpp @@ -27,6 +27,7 @@ #include #include #include +#include extern "C" { #include "librocket.h" @@ -41,9 +42,10 @@ extern "C" { #define ROCKET_TILE_N 128 #define ROCKET_K_MAX 8192 #define ROCKET_MIN_BATCH 32 // below this, CPU wins (decode is bandwidth-bound) +#define ROCKET_NFD 3 // one fd/entity caps at 2 NPU cores (DRM sched tie-break); use 3 for all cores struct ggml_backend_rocket_context { - int fd = -1; + int fd[ROCKET_NFD] = { -1, -1, -1 }; // reusable scratch (grown as needed, never shrunk) std::vector wf; // dequantized weight plane [N*K] std::vector aq; // quantized activations [M*K] @@ -60,6 +62,20 @@ static inline uint8_t rocket_q8(float v, float inv_s) { return (uint8_t)(q + 128); // symmetric int8 stored as uint8, zp=128 } +struct rkt_thread_arg { + int fd; const uint8_t *X, *Wc; uint32_t M, N, K; + float in_scale, out_scale; const float *out_scales; + uint32_t col_start, col_num; uint8_t *Y; int rc; +}; +static void * rkt_thread_fn(void *a) { + struct rkt_thread_arg *t = (struct rkt_thread_arg *)a; + t->rc = rkt_npu_matmul(t->fd, t->X, t->Wc, NULL, t->M, t->N, t->K, + 128, 128, 128, t->in_scale, 1.0f, t->out_scale, + t->out_scales, ROCKET_TILE_M, ROCKET_TILE_N, + t->col_start, t->col_num, t->Y); + return NULL; +} + static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; // weight [K, N] const struct ggml_tensor * src1 = dst->src[1]; // activ [K, M] @@ -168,16 +184,29 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc } } - static const bool dbg = getenv("GGML_ROCKET_DEBUG") != NULL; - if (dbg) fprintf(stderr, "[rocket] mul_mat M=%lld N=%lld K=%lld plane(%lld,%lld) type=%s ...", - (long long)M, (long long)N, (long long)K, - (long long)i12, (long long)i13, ggml_type_name(type)), fflush(stderr); - int rc = rkt_npu_matmul(ctx->fd, ctx->aq.data(), ctx->wq.data(), NULL, - (uint32_t)M, (uint32_t)N, (uint32_t)K, - 128, 128, 128, - a_scale, 1.0f, out_scale, ctx->os.data(), - ROCKET_TILE_M, ROCKET_TILE_N, ctx->yq.data()); - if (dbg) fprintf(stderr, " rc=%d\n", rc); + // Split N columns across ROCKET_NFD fds == NPU cores (a single fd + // caps at 2 cores via the DRM sched tie-break). Column ranges are + // TILE_N-aligned so the per-tile weight-scale blocks don't split. + pthread_t th[ROCKET_NFD]; + struct rkt_thread_arg ta[ROCKET_NFD]; + uint32_t per = ((uint32_t)N + ROCKET_NFD - 1) / ROCKET_NFD; + per = ((per + ROCKET_TILE_N - 1) / ROCKET_TILE_N) * ROCKET_TILE_N; + int nth = 0; + for (int fdi = 0; fdi < ROCKET_NFD; fdi++) { + uint32_t c0 = (uint32_t)fdi * per; + if (c0 >= (uint32_t)N) break; + uint32_t cn = (c0 + per <= (uint32_t)N) ? per : ((uint32_t)N - c0); + ta[nth] = (struct rkt_thread_arg){ ctx->fd[fdi], ctx->aq.data(), + ctx->wq.data(), (uint32_t)M, (uint32_t)N, (uint32_t)K, + a_scale, out_scale, ctx->os.data(), c0, cn, ctx->yq.data(), 0 }; + pthread_create(&th[nth], NULL, rkt_thread_fn, &ta[nth]); + nth++; + } + int rc = 0; + for (int i = 0; i < nth; i++) { + pthread_join(th[i], NULL); + if (ta[i].rc) rc = ta[i].rc; + } GGML_ASSERT(rc == 0 && "rkt_npu_matmul failed"); // --- dequantize int8 result into dst (F32) --- @@ -260,14 +289,16 @@ static ggml_guid_t ggml_backend_rocket_guid(void) { } ggml_backend_t ggml_backend_rocket_init(void) { - int fd = rocket_open("/dev/accel/accel0"); - if (fd < 0) { - GGML_LOG_ERROR("%s: failed to open /dev/accel/accel0 (%d)\n", __func__, fd); - return NULL; - } - ggml_backend_rocket_context * ctx = new ggml_backend_rocket_context; - ctx->fd = fd; + for (int i = 0; i < ROCKET_NFD; i++) { + int fd = rocket_open("/dev/accel/accel0"); + if (fd < 0) { + GGML_LOG_ERROR("%s: failed to open /dev/accel/accel0 (%d)\n", __func__, fd); + delete ctx; + return NULL; + } + ctx->fd[i] = fd; + } ggml_backend_t backend = new ggml_backend { /* .guid = */ ggml_backend_rocket_guid(), diff --git a/ggml/src/ggml-rocket/rkt_gemm.c b/ggml/src/ggml-rocket/rkt_gemm.c index b00c70c5f..40e88ef61 100644 --- a/ggml/src/ggml-rocket/rkt_gemm.c +++ b/ggml/src/ggml-rocket/rkt_gemm.c @@ -6,22 +6,25 @@ #include "rkt_matmul.h" int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n, + uint32_t col_start, uint32_t col_num, struct rkt_gemm_tile *tiles, int max_tiles) { if (M == 0 || N == 0 || tile_m == 0 || tile_n == 0) return -1; + uint32_t col_end = col_start + col_num; + if (col_end > N) col_end = N; int t = 0; - /* column-major: consecutive tiles share the column tile, so the weight - * pack+write can be cached across all row tiles of a column. */ - for (uint32_t c = 0; c < N; c += tile_n) { + /* column-major over [col_start,col_end): consecutive tiles share the column + * tile (weight pack+write cached across a column's row tiles). */ + for (uint32_t c = col_start; c < col_end; c += tile_n) { for (uint32_t r = 0; r < M; r += tile_m) { if (t >= max_tiles) return -1; tiles[t].row = r; tiles[t].col = c; tiles[t].m = (r + tile_m <= M) ? tile_m : (M - r); - tiles[t].n = (c + tile_n <= N) ? tile_n : (N - c); + tiles[t].n = (c + tile_n <= col_end) ? tile_n : (col_end - c); t++; } } diff --git a/ggml/src/ggml-rocket/rkt_gemm.h b/ggml/src/ggml-rocket/rkt_gemm.h index dd4bccc90..a1012c8d1 100644 --- a/ggml/src/ggml-rocket/rkt_gemm.h +++ b/ggml/src/ggml-rocket/rkt_gemm.h @@ -28,6 +28,7 @@ struct rkt_gemm_tile { * rkt_gemm_op_fits() to confirm a tile actually fits one rocket op. */ int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n, + uint32_t col_start, uint32_t col_num, struct rkt_gemm_tile *tiles, int max_tiles); /* True (1) if an m x n, full-k matmul fits a single rocket op. */ diff --git a/ggml/src/ggml-rocket/rkt_matmul.c b/ggml/src/ggml-rocket/rkt_matmul.c index b61d9ac6b..4a9c4d7b4 100644 --- a/ggml/src/ggml-rocket/rkt_matmul.c +++ b/ggml/src/ggml-rocket/rkt_matmul.c @@ -83,8 +83,6 @@ static unsigned calc_weights_banks(const struct op *o) return banks; } -int rkt_g_task_num = 0; /* experiment: CBUF weight-reuse multi-task chaining */ - int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, uint32_t M, uint32_t N, uint32_t K, uint64_t input_dma, uint64_t weights_dma, @@ -93,7 +91,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, int32_t weight_zero_point, int32_t output_zero_point, float input_scale, float weights_scale, - float output_scale, uint64_t bias_dma) + float output_scale, uint64_t bias_dma, int task_num) { if (M == 0 || N == 0 || K == 0) return -1; @@ -210,7 +208,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, c.reuse_weights_cbuf = reuse_weights_cbuf; c.addition_input = o.addition_input; c.add_tensor = o.add_tensor; - c.task_num = rkt_g_task_num; + c.task_num = task_num; /* ---- CORE/DPU/PC params ---- */ struct coredpu_params d; @@ -260,5 +258,5 @@ int rkt_build_matmul_regcmd(uint64_t *out, int out_capacity, input_dma, weights_dma, output_dma, input_zero_point, weight_zero_point, output_zero_point, - 1.0f, 1.0f, 1.0f, 0); + 1.0f, 1.0f, 1.0f, 0, 0); } diff --git a/ggml/src/ggml-rocket/rkt_matmul.h b/ggml/src/ggml-rocket/rkt_matmul.h index d5c58b72e..4ebc74d7e 100644 --- a/ggml/src/ggml-rocket/rkt_matmul.h +++ b/ggml/src/ggml-rocket/rkt_matmul.h @@ -51,7 +51,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, int32_t input_zero_point, int32_t weight_zero_point, int32_t output_zero_point, float input_scale, float weights_scale, - float output_scale, uint64_t bias_dma); + float output_scale, uint64_t bias_dma, int task_num); int rkt_build_matmul_regcmd(uint64_t *out, int out_capacity, uint32_t M, uint32_t N, uint32_t K, diff --git a/ggml/src/ggml-rocket/rkt_npu_matmul.c b/ggml/src/ggml-rocket/rkt_npu_matmul.c index d0b4fd4c8..f620aa5b9 100644 --- a/ggml/src/ggml-rocket/rkt_npu_matmul.c +++ b/ggml/src/ggml-rocket/rkt_npu_matmul.c @@ -38,8 +38,6 @@ __attribute__((destructor)) static void rkt_prof_dump(void){ g_setup_ns/1e9,g_submit_ns/1e9,g_prep_ns/1e9,g_free_ns/1e9); } -extern int rkt_g_task_num; - int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, uint32_t M, uint32_t N, uint32_t K, uint64_t input_dma, uint64_t weights_dma, @@ -47,7 +45,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity, int32_t input_zero_point, int32_t weight_zero_point, int32_t output_zero_point, float input_scale, float weights_scale, - float output_scale, uint64_t bias_dma); + float output_scale, uint64_t bias_dma, int task_num); struct rkt_bo { uint32_t h; uint64_t dma, off; void *map; uint32_t sz; }; @@ -87,7 +85,8 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc, const int32_t *bias, uint32_t M, uint32_t N, uint32_t K, uint8_t izp, uint8_t wzp, uint8_t ozp, float in_scale, float wt_scale, float out_scale, const float *out_scales, - uint32_t tile_m, uint32_t tile_n, uint8_t *Y) + uint32_t tile_m, uint32_t tile_n, + uint32_t col_start, uint32_t col_num, uint8_t *Y) { int ret = -1; int32_t *bias0 = NULL; @@ -110,15 +109,15 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc, uint8_t *wpk = calloc(1, max_wsz); int32_t *bpk = malloc((size_t)tile_n * sizeof(int32_t)); struct rkt_gemm_tile *tiles = malloc( - ((M + tile_m - 1) / tile_m) * ((N + tile_n - 1) / tile_n) * + ((M + tile_m - 1) / tile_m) * ((col_num + tile_n - 1) / tile_n) * sizeof(struct rkt_gemm_tile)); uint64_t *rc = malloc(4096 * sizeof(uint64_t)); if (!Xt || !Wt || !ipk || !wpk || !bpk || !tiles || !rc) goto out; - int nt = rkt_gemm_plan(M, N, tile_m, tile_n, tiles, + int nt = rkt_gemm_plan(M, N, tile_m, tile_n, col_start, col_num, tiles, (int)(((M + tile_m - 1) / tile_m) * - ((N + tile_n - 1) / tile_n))); + ((col_num + tile_n - 1) / tile_n))); if (nt <= 0) goto out; @@ -168,12 +167,12 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc, bo_write(fd, &rin[cnt], ipk, rkt_raw_input_size(1, m, K)); bo_write(fd, &ro[cnt], NULL, rkt_raw_output_size(1, m, n)); - rkt_g_task_num = cnt; /* task 0 loads weights, 1+ reuse CBUF */ + /* task 0 loads weights DDR->CBUF, tasks 1+ reuse SRAM */ int nw = rkt_build_matmul_regcmd_scaled( rc, 4096, m, n, K, rin[cnt].dma, w.dma, ro[cnt].dma, izp, wzp, ozp, in_scale, wt_scale, - out_scales ? out_scales[cc] : out_scale, b.dma); - if (nw < 0) { rkt_g_task_num = 0; goto out; } + out_scales ? out_scales[cc] : out_scale, b.dma, cnt); + if (nw < 0) goto out; bo_write(fd, &rreg[cnt], rc, (unsigned)nw * sizeof(uint64_t)); tasks[cnt].regcmd = (uint32_t)rreg[cnt].dma; @@ -182,7 +181,6 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc, meta[cnt].r = r; meta[cnt].m = m; meta[cnt].n = n; cnt++; t++; } - rkt_g_task_num = 0; struct drm_rocket_job job; memset(&job, 0, sizeof job); diff --git a/ggml/src/ggml-rocket/rkt_npu_matmul.h b/ggml/src/ggml-rocket/rkt_npu_matmul.h index 705237bc0..5fd5fbfba 100644 --- a/ggml/src/ggml-rocket/rkt_npu_matmul.h +++ b/ggml/src/ggml-rocket/rkt_npu_matmul.h @@ -30,6 +30,7 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc, const int32_t *bias, uint32_t M, uint32_t N, uint32_t K, uint8_t izp, uint8_t wzp, uint8_t ozp, float in_scale, float wt_scale, float out_scale, const float *out_scales, - uint32_t tile_m, uint32_t tile_n, uint8_t *Y); + uint32_t tile_m, uint32_t tile_n, + uint32_t col_start, uint32_t col_num, uint8_t *Y); #endif /* RKT_NPU_MATMUL_H */