Просмотр исходного кода

ggml : fix soft max out-of-bounds access (#4307)

ggml-ci
Georgi Gerganov 2 лет назад
Родитель
Сommit
adf3de4f69
1 измененных файлов с 4 добавлено и 1 удалено
  1. 4 1
      ggml.c

+ 4 - 1
ggml.c

@@ -15629,7 +15629,6 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             } break;
             } break;
         case GGML_OP_DIAG_MASK_ZERO:
         case GGML_OP_DIAG_MASK_ZERO:
         case GGML_OP_DIAG_MASK_INF:
         case GGML_OP_DIAG_MASK_INF:
-        case GGML_OP_SOFT_MAX:
         case GGML_OP_SOFT_MAX_BACK:
         case GGML_OP_SOFT_MAX_BACK:
         case GGML_OP_ROPE:
         case GGML_OP_ROPE:
         case GGML_OP_ROPE_BACK:
         case GGML_OP_ROPE_BACK:
@@ -15645,6 +15644,10 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             {
             {
                 n_tasks = 1; //TODO
                 n_tasks = 1; //TODO
             } break;
             } break;
+        case GGML_OP_SOFT_MAX:
+            {
+                n_tasks = MIN(MIN(4, n_threads), ggml_nrows(node->src[0]));
+            } break;
         case GGML_OP_CONV_TRANSPOSE_1D:
         case GGML_OP_CONV_TRANSPOSE_1D:
             {
             {
                 n_tasks = n_threads;
                 n_tasks = n_threads;