Merge pull request #3 from lanhin/tensor_export_turnoff

Tensor export turnoff

Author: lanhin

dpu/dpu_main.c

@@ -11,16 +11,26 @@
 #include <alloc.h>
 #include <barrier.h>
 #include <seqread.h>
+#include <mutex_pool.h>
 
 #define PIM_KERNEL_DPU 1
 #include "../ggml/include/ggml.h"
 #define GGML_COMMON_DECL_C
 #include "../ggml/src/ggml-common.h"
 
 #define PRINT 0
+#define SEGMENT_PER_ROW 4
+
+// Find the lowest index for the rank-th group
+#define BLOCK_LOW(rank, size, n) ((rank) * (n) / (size))
+
+// Find the highest index for the rank-th group
+#define BLOCK_HIGH(rank, size, n) (BLOCK_LOW((rank) + 1, (size), (n)) - 1)
 
 __mram_ptr float *ptable_f32_f16;
 
+__host int16_t mul_table_int4_int8[1<<4][1<<8];
+
 inline static float lookup_fp16_to_fp32(uint16_t f) {
     uint16_t s;
     memcpy(&s, &f, sizeof(uint16_t));
@@ -35,6 +45,7 @@ inline static float lookup_fp16_to_fp32(uint16_t f) {
 
 // Barrier
 BARRIER_INIT(my_barrier, NR_TASKLETS);
+MUTEX_POOL_INIT(g_psumf_mutex_pool, NR_TASKLETS);
 
 /*
 DPU MRAM Memory:
@@ -91,8 +102,9 @@ int wram2mram(__mram_ptr void *pmram,void *pwram,uint32_t size)
 }
 
 
-// set psumf to global value for each thread access
-static float *psumf = NULL;
+// set g_psumf to global value for each thread access
+static float *g_psumf = NULL;
+static block_q8_0 *g_pinput_cache = NULL;
 
 void init(unsigned int tasklet_id) {
 #if PRINT
@@ -140,9 +152,11 @@ int main() {
 #endif
 
     // set sart line, end line and line number in each thread
-    uint16_t weight_rows_per_thread = cache_meta->rows_per_dpu / NR_TASKLETS;
-    uint16_t weight_start_row = tasklet_id * weight_rows_per_thread;
-    uint16_t weight_end_row = weight_start_row + weight_rows_per_thread;
+    uint16_t segments_num = cache_meta->rows_per_dpu * SEGMENT_PER_ROW;
+    uint16_t segment_start = BLOCK_LOW(tasklet_id, NR_TASKLETS, segments_num);
+    uint16_t segment_end = BLOCK_HIGH(tasklet_id, NR_TASKLETS, segments_num);
+
+    assert(segment_start <= segment_end && "There are not enough segments to allocate to the tasklets");
 
     // todo:rest row is existed, first thread in every dpu can one more row
     uint16_t weight_rows_cur_thread;
@@ -184,83 +198,82 @@ int main() {
             return -1;
         }
         int nb = pinputcache->ne[0]/QK8_0;
+
+        assert(SEGMENT_PER_ROW <= nb && nb % SEGMENT_PER_ROW == 0 
+            && "Too many segments are allocated to each row.");
+
         int qk = QK8_0;
         input_row_size = nb*sizeof(block_q8_0);
         __mram_ptr void *pweight_base = (__mram_ptr void *)(weightmetadatabase + sizeof(struct pim_meta));
         __mram_ptr void *pinput_base = DPU_MRAM_HEAP_POINTER + cache_meta->input_offset + sizeof(pim_matrix_des);
-
+        
         if (tasklet_id == 0) {
-            psumf = (float *)mem_alloc(sizeof(float)*input_cols*weight_rows_cur_thread);
+            g_psumf = (float *)mem_alloc(sizeof(float)*input_cols*weight_rows_cur_thread);
+            g_pinput_cache = (block_q8_0 *) mem_alloc(sizeof(block_q8_0) * nb);
+            memset(g_psumf, 0 ,sizeof(float)*input_cols*weight_rows_cur_thread);
         }
-        barrier_wait(&my_barrier);
 
-        // psumf = (float *)mem_alloc(sizeof(float)*input_cols*weight_rows_cur_thread);
-        memset(psumf, 0 ,sizeof(float)*input_cols*weight_rows_cur_thread);
-        
 #if PRINT
         printf("input_cols=%d, rows_cur_thread=%d, nb=%d, input_row_size=%d\n",input_cols,weight_rows_cur_thread,nb,input_row_size);
 #endif
-        block_q4_0 *pweight_cache = (block_q4_0 *) mem_alloc(sizeof(block_q4_0)*nb);
-        block_q8_0 *pinput_cache = (block_q8_0 *) mem_alloc(sizeof(block_q8_0)*nb);
+
+        uint16_t segment_nb_size = nb / SEGMENT_PER_ROW;
+        block_q4_0 *pweight_cache = (block_q4_0 *) mem_alloc(sizeof(block_q4_0) * segment_nb_size);
 
         // weight_rows_cur_thread = 16;
         for(int l = 0;l < input_cols;l++) {
-          __mram_ptr block_q8_0 *pinput = pinput_base + l * nb * sizeof(block_q8_0);
-            mram2wram(pinput, pinput_cache, sizeof(block_q8_0)*nb);
-#if PRINT
-            printf("input:\n");
-            for (int i = 0; i < nb; i++) {
-              printf("d=%u\n",pinput[i].d);
-              for (int kkk=0;kkk<QK8_0;kkk++) {
-                printf("%d ",pinput[i].qs[kkk]);
-              }
-            printf("\n");
+            if (tasklet_id == 0) {
+                __mram_ptr block_q8_0 *pinput = pinput_base + l * nb * sizeof(block_q8_0);
+                mram2wram(pinput, g_pinput_cache, sizeof(block_q8_0)*nb);
             }
-            printf("pweight_base: %p\n", pweight_base);
-#endif
-            // for(int k = 0;k < weight_rows_cur_thread;k++) {
-            for (int k = weight_start_row; k < weight_end_row; ++k) {
-              __mram_ptr block_q4_0 *pweight = pweight_base + pinputcache->layerid * cache_meta->layer_len + k * nb * sizeof(block_q4_0);
-                mram2wram(pweight, pweight_cache, sizeof(block_q4_0)*nb);
-#if PRINT
-                if (k % 64 == 0) {
-                  printf("pweight_cache[%d].d=%d\n pweight_cache[%d].qs=", k*128, pweight_cache[0].d, k*128);
-                  for (int kkk=0;kkk<QK4_0/2;kkk++) {
-                    int v0 = (pweight_cache[0].qs[kkk] & 0x0f) - 8;
-                    int v1 = (pweight_cache[0].qs[kkk]  >> 4) - 8;
-                    printf(" %d, %d", v0, v1);
-                  }
-                  printf("\n");
-                }
-#endif
 
-                for (int i = 0; i < nb; i++) {
-                    //printf("input_col:%d, current inner weight row idx:%d\n",l,k);
+            barrier_wait(&my_barrier);
+
+            __mram_ptr block_q4_0 *pweight_addr = pweight_base + pinputcache->layerid * cache_meta->layer_len;
 
+            for (int k = segment_start; k <= segment_end; ++k) {
+                __mram_ptr block_q4_0 *pweight = pweight_addr + k * segment_nb_size;
+                mram2wram(pweight, pweight_cache, sizeof(block_q4_0) * segment_nb_size);
+
+                block_q8_0 *pinput_cache = g_pinput_cache + k % SEGMENT_PER_ROW * segment_nb_size;
+
+                for (int i = 0; i < segment_nb_size; i++) {
                     int sumi = 0;
                     for (int j = 0; j < qk/2; ++j) {
-                        const int v0 = (pweight_cache[i].qs[j] & 0x0F) - 8;
-                        const int v1 = (pweight_cache[i].qs[j] >>   4) - 8;
+                        const int8_t v0 = (pweight_cache[i].qs[j] & 0x0F) - 8;
+                        const int8_t v1 = (pweight_cache[i].qs[j] >>   4) - 8;
 
-                        sumi += (v0 * pinput_cache[i].qs[j]) + (v1 * pinput_cache[i].qs[j + qk/2]);
+                        // sumi += (v0 * pinput_cache[i].qs[j]) + (v1 * pinput_cache[i].qs[j + qk/2]);
+                        sumi += mul_table_int4_int8[v0 + 8][pinput_cache[i].qs[j] - INT8_MIN] + 
+                                mul_table_int4_int8[v1 + 8][pinput_cache[i].qs[j + qk/2] - INT8_MIN];
                     }
-
-                    psumf[l*weight_rows_cur_thread + k] += sumi*FP16_TO_FP32(pweight_cache[i].d)*FP16_TO_FP32(pinput_cache[i].d);
+                    
+                    int psumf_idx = l * weight_rows_cur_thread + k / SEGMENT_PER_ROW;
+                    float sum = sumi * FP16_TO_FP32(pweight_cache[i].d) * FP16_TO_FP32(pinput_cache[i].d);
+                    mutex_pool_lock(&g_psumf_mutex_pool, psumf_idx);
+                    g_psumf[psumf_idx] += sum;
+                    // g_psumf[psumf_idx] += sumi;
+                    mutex_pool_unlock(&g_psumf_mutex_pool, psumf_idx);
                 }
             }
         }
     }
 
-    offset += (sizeof(pim_matrix_des) + input_row_size * input_cols);
-#if PRINT
-    for(int iii=0;iii<cache_meta->rows_per_dpu;iii+=128) {
-        printf("psumf[%d]=%f\n",iii,psumf[iii]);
+    barrier_wait(&my_barrier);
+
+    if (tasklet_id == 0){
+        offset += (sizeof(pim_matrix_des) + input_row_size * input_cols);
+        #if PRINT
+            for(int iii=0;iii<cache_meta->rows_per_dpu;iii+=128) {
+                printf("g_psumf[%d]=%f\n",iii,g_psumf[iii]);
+            }
+        
+            printf("output offset=%d\n",offset);
+        #endif
+        // Write C Matrix to current MRAM block
+        // Note: with input_cols > 1, the results should be rearranged on host
+        wram2mram((__mram_ptr void *) (DPU_MRAM_HEAP_POINTER + offset), g_psumf, sizeof(float)*input_cols*weight_rows_cur_thread);
     }
 
-    printf("output offset=%d\n",offset);
-#endif
-    // Write C Matrix to current MRAM block
-    // Note: with input_cols > 1, the results should be rearranged on host
-    wram2mram((__mram_ptr void *) (DPU_MRAM_HEAP_POINTER + offset), psumf, sizeof(float)*input_cols*weight_rows_cur_thread);
     return 0;
 }

dpu/pim_build.sh

@@ -1,2 +1,2 @@
 #!/bin/bash
-dpu-upmem-dpurte-clang -Wall -Wextra -O2 -DNR_TASKLETS=8 -DBL=11 -o gemv_dpu dpu_main.c
+dpu-upmem-dpurte-clang -Wall -Wextra -O3 -DNR_TASKLETS=16 -DBL=11 -o gemv_dpu dpu_main.c

examples/tensor/ts.cpp

@@ -3,12 +3,12 @@
 #include <iomanip>
 #include <chrono>
 
-#include <vector>
-
-#define NR_DPUS 2048
+#define NR_DPUS 512
 #define NR_LAYER 2
 #define DPU_BINARY "./dpu/gemv_dpu"
 
+int16_t mul_table_int4_int8[1<<4][1<<8];
+
 void fp_table_init(void) {
   for (int i = 0; i < (1 << 16); ++i) {
                 union {
@@ -19,12 +19,22 @@ void fp_table_init(void) {
             }
 }
 
+void mul_table_int4_int8_init(void) {
+  for(int i = 0; i < (1 << 4); ++i){
+    for(int j = 0; j< (1 << 8); ++j){
+      mul_table_int4_int8[i][j] = (i - 8) * (j + INT8_MIN);
+    }
+  }
+}
+
+#ifdef PIM_KERNEL
 int gemv_dpu_kernel(struct pim_context *context, struct ggml_tensor * w, struct ggml_tensor * in_q, struct ggml_tensor * res) {
   uint32_t pim_offset = 0;
   struct dpu_set_t dpu;
 
   std::chrono::high_resolution_clock::time_point ex_tp1 = std::chrono::high_resolution_clock::now();
 
+  DPU_ASSERT(dpu_broadcast_to(context->dpu_set, "mul_table_int4_int8", 0, (void *)(mul_table_int4_int8), sizeof(mul_table_int4_int8), DPU_XFER_DEFAULT));
   //ggml_table_f32_f16 tbl is transferred to pim
   DPU_ASSERT(dpu_broadcast_to(context->dpu_set, DPU_MRAM_HEAP_POINTER_NAME, pim_offset, (void *)(ggml_table_f32_f16), sizeof(ggml_table_f32_f16), DPU_XFER_DEFAULT));
   pim_offset += sizeof(ggml_table_f32_f16);
@@ -104,8 +114,8 @@ int gemv_dpu_kernel(struct pim_context *context, struct ggml_tensor * w, struct
 
   dur = ex_tp2 - ex_tp1;
 
-  std::cout << "dpu: 执行用时：" << std::chrono::duration_cast<std::chrono::milliseconds>(dur).count() << " ms" << std::endl;
-  std::cout << "dpu: 执行用时：" << std::chrono::duration_cast<std::chrono::microseconds>(dur).count() << " us" << std::endl;
+  // std::cout << "执行用时：" << std::chrono::duration_cast<std::chrono::milliseconds>(dur).count() << " ms" << std::endl;
+  std::cout << "执行用时：" << std::chrono::duration_cast<std::chrono::microseconds>(dur).count() << " us" << std::endl;
 
   // Check results
   float *mul_mat_res = (float *)res->data;
@@ -116,6 +126,7 @@ int gemv_dpu_kernel(struct pim_context *context, struct ggml_tensor * w, struct
 
   return 0;
 }
+#endif
 
 
 void gemv_cpu_kernel(struct pim_context *context, struct ggml_tensor * w, struct ggml_tensor * in_q, struct ggml_tensor * res_comp) {
@@ -163,7 +174,9 @@ void gemv_cpu_kernel(struct pim_context *context, struct ggml_tensor * w, struct
 int main(int argc, char** argv) {
   // init fp table for fp16 dump
   fp_table_init();
+  mul_table_int4_int8_init();
 
+#ifdef PIM_KERNEL
   // WQ-PIM allocate dpu
   struct pim_context *pqcontext = (struct pim_context *)malloc(sizeof(struct pim_context));
   memset(pqcontext,0,sizeof(struct pim_context));
@@ -213,6 +226,6 @@ int main(int argc, char** argv) {
   // float first_res = mul_add_q4_0_q8_0(ts_a, ts_bq);
   // std::cout<<"first element: "<<std::fixed << std::setprecision(6)<<first_res<<std::endl;
 
-  
+#endif
   return 0;
 }