forward - OptiTrust Trace

Trace for forward

(Function.inline [ f; cCall "forward" ]);
(Loop.tile ~bound:TileBoundMin (trm_int chunk_len) [ f; cFor "i" ]);
Loop.hoist [nbMulti;cFunDef "generate_prompt_proc"; cVarDefs [ "embedding"; "mha_norm"; "mha_q"; "mha_score"; "mha_att"; "mha_blend"; "mha_out"; "ffn_norm"; "ffn_up"; "ffn_fc"; "ffn_out" ]; ];
Loop.fission [ f; cForBody "i"; tBetweenAll ];
Loop.reorder_at ~order:[ "l"; "i" ] [ f; cForBody "l"; dSeqNth 0 ];
Loop.fission [ f; cForBody "l"; cForBody "i"; tBetweenAll ];
Loop.reorder_at ~order:[ "q"; "i" ] [ nbMulti; f; cForBody "q"; dSeqNth 0 ];
Loop.reorder_at ~order:[ "h"; "i" ] [ nbMulti; f; cForBody "h"; dSeqNth 0 ];
Matrix.reorder_dims ~order:[ 1; 0; 2 ] [ nbMulti; f; cVarDefs [ "mha_q"; "mha_score"; "mha_blend" ] ];
Matrix.tile ~block_size:q_head_per_kv_head_count ~nb_blocks:kv_headcount ~index_dim:0 [ nbMulti; f; cVarDefs [ "mha_q"; "mha_score"; "mha_blend" ] ];
Loop_basic.grid_enumerate [ ("h2", kv_headcount); ("q2", q_head_per_kv_head_count) ] [ nbMulti; f; cFor "q" ];
Variable.inline [ nbMulti; f; cVarDef "q" ];
Rewrite.equiv_at "int i; int j ; int k; ==> (i*j +k) /j == i" ~indepth:true [ f ];
Rewrite.equiv_at "int i; int j ; int k; ==> (i*j +k) %j == k" ~indepth:true [ f ];
Function.inline [ nbMulti; f; cCall "matvec" ];
Matrix.simpl_access_of_access ~indepth:true [ f ];
Matrix.simpl_index_add [ nbMulti; f; cCellAccess ~base:[ cVar ~substr:true "" ] (); cBinop ~lhs:[ cCall ~regexp:true "MINDEX." ] Binop_add ];
Rewrite.equiv_at "int j; ==> 0 + j == j" [ nbMulti; f ] ~indepth:true;
Function.uninline ~f:[ cFunDef "matmul" ] [ occIndices [ 0; 3; 4; 5 ]; f; cFor "i" ~body:[ cSeq ~instrs_pred:(target_list_one_st [ cFor "j" ]) () ] ])

advancedargumentsjustification

tmp/{before6fc5cf.cpp → after016d03.cpp} RENAMED Viewed

@@ -184,20 +184,238 @@ void generate_prompt_proc(int vocabulary_len, int context_len, int layer_count,
              int q_head_per_kv_head_count, int embedding_dim,
              int head_dim, int q_dim, int kv_dim, int hidden_dim,
              float epsilon, float* embedding_weight,
              float* mha_norm_weight, float* mha_q_weight,
              float* mha_k_weight, float* mha_v_weight,
              float* mha_out_weight, float* ffn_norm_weight,
              float* ffn_fc_weight, float* ffn_up_weight,
              float* ffn_out_weight, float* out_norm_weight,
              float* out_weight, float* k_cache, float* v_cache,
              float* logits, int* sequence, int sequence_len) {
- for (int i = 0; i < sequence_len; i++) {
-  forward(sequence[i], vocabulary_len, context_len, layer_count, q_head_count,
-      kv_head_count, q_head_per_kv_head_count, embedding_dim, head_dim,
-      q_dim, kv_dim, hidden_dim, epsilon, embedding_weight,
-      mha_norm_weight, mha_q_weight, mha_k_weight, mha_v_weight,
-      mha_out_weight, ffn_norm_weight, ffn_fc_weight, ffn_up_weight,
-      ffn_out_weight, out_norm_weight, out_weight, k_cache, v_cache,
-      logits, i, i == sequence_len - 1 ? 1 : 0);
  }
 }

              int q_head_per_kv_head_count, int embedding_dim,
              int head_dim, int q_dim, int kv_dim, int hidden_dim,
              float epsilon, float* embedding_weight,
              float* mha_norm_weight, float* mha_q_weight,
              float* mha_k_weight, float* mha_v_weight,
              float* mha_out_weight, float* ffn_norm_weight,
              float* ffn_fc_weight, float* ffn_up_weight,
              float* ffn_out_weight, float* out_norm_weight,
              float* out_weight, float* k_cache, float* v_cache,
              float* logits, int* sequence, int sequence_len) {
+ for (int bi = 0; bi < sequence_len; bi += 512) {
+  float* const embedding =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  float* const mha_norm =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  float* const mha_q =
+    (float*)malloc(MSIZE4(kv_head_count, q_head_per_kv_head_count,
+               min(sequence_len, bi + 512) - bi, head_dim) *
+            sizeof(float));
+  __ghost([&]() {
+   __consumes(
+     "for i1 in 0..q_head_count -> for i2 in 0..(min(sequence_len, bi + "
+     "512) - bi) -> for i3 in 0..head_dim -> &mha_q[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __produces(
+     "for i1 in 0..(min(sequence_len, bi + 512) - bi) -> for i2 in "
+     "0..q_head_count -> for i3 in 0..head_dim -> &mha_q[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __admitted();
+   __with("justif := reorder_groups");
+  });
+  float* const mha_score =
+    (float*)malloc(MSIZE4(kv_head_count, q_head_per_kv_head_count,
+               min(sequence_len, bi + 512) - bi, context_len) *
+            sizeof(float));
+  __ghost([&]() {
+   __consumes(
+     "for i1 in 0..q_head_count -> for i2 in 0..(min(sequence_len, bi + "
+     "512) - bi) -> for i3 in 0..context_len -> &mha_score[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __produces(
+     "for i1 in 0..(min(sequence_len, bi + 512) - bi) -> for i2 in "
+     "0..q_head_count -> for i3 in 0..context_len -> &mha_score[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __admitted();
+   __with("justif := reorder_groups");
+  });
+  float* const mha_blend =
+    (float*)malloc(MSIZE4(kv_head_count, q_head_per_kv_head_count,
+               min(sequence_len, bi + 512) - bi, head_dim) *
+            sizeof(float));
+  __ghost([&]() {
+   __consumes(
+     "for i1 in 0..q_head_count -> for i2 in 0..(min(sequence_len, bi + "
+     "512) - bi) -> for i3 in 0..head_dim -> &mha_blend[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __produces(
+     "for i1 in 0..(min(sequence_len, bi + 512) - bi) -> for i2 in "
+     "0..q_head_count -> for i3 in 0..head_dim -> &mha_blend[i1 / "
+     "q_head_per_kv_head_count][i1 % q_head_per_kv_head_count][i2][i3] ~> "
+     "UninitCell");
+   __admitted();
+   __with("justif := reorder_groups");
+  });
+  float* const mha_att =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  float* const mha_out =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  float* const ffn_norm =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  float* const ffn_fc = (float*)malloc(
+    MSIZE2(min(sequence_len, bi + 512) - bi, hidden_dim) * sizeof(float));
+  float* const ffn_up = (float*)malloc(
+    MSIZE2(min(sequence_len, bi + 512) - bi, hidden_dim) * sizeof(float));
+  float* const ffn_out =
+    (float*)malloc(MSIZE2(min(sequence_len, bi + 512) - bi, embedding_dim) *
+            sizeof(float));
+  for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+   __ghost(assume, "P := in_range(i + bi, bi..min(sequence_len, bi + 512))");
+  }
+  for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+   for (int e = 0; e < embedding_dim; e++) {
+    embedding[i][e] = embedding_weight[sequence[i + bi]][e];
+   }
+  }
+  for (int l = 0; l < layer_count; l++) {
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    rmsnorm(embedding_dim, &mha_norm[i][0], &embedding[i][0],
+        &mha_norm_weight[l][0], epsilon);
+   }
+   for (int h2 = 0; h2 < kv_head_count; h2++) {
+    for (int q2 = 0; q2 < q_head_per_kv_head_count; q2++) {
+     matmul(min(sequence_len, bi + 512) - bi, head_dim, embedding_dim,
+         &mha_q[h2][q2][0][0], mha_norm,
+         &mha_q_weight[l][h2 * q_head_per_kv_head_count + q2][0][0]);
+    }
+   }
+   for (int h = 0; h < kv_head_count; h++) {
+    for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+     for (int j = 0; j < head_dim; j++) {
+      k_cache[l][h][i + bi][j] = 0.f;
+      for (int k = 0; k < embedding_dim; k++) {
+       k_cache[l][h][i + bi][j] +=
+         mha_norm[i][k] * mha_k_weight[l][h][j][k];
+      }
+     }
+    }
+   }
+   for (int h = 0; h < kv_head_count; h++) {
+    for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+     for (int j = 0; j < head_dim; j++) {
+      v_cache[l][h][i + bi][j] = 0.f;
+      for (int k = 0; k < embedding_dim; k++) {
+       v_cache[l][h][i + bi][j] +=
+         mha_norm[i][k] * mha_v_weight[l][h][j][k];
+      }
+     }
+    }
+   }
+   for (int h2 = 0; h2 < kv_head_count; h2++) {
+    for (int q2 = 0; q2 < q_head_per_kv_head_count; q2++) {
+     for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+      rope(head_dim, &mha_q[h2][q2][i][0], i + bi);
+     }
+    }
+   }
+   for (int h = 0; h < kv_head_count; h++) {
+    for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+     rope(head_dim, &k_cache[l][h][i + bi][0], i + bi);
+    }
+   }
+   for (int h2 = 0; h2 < kv_head_count; h2++) {
+    for (int q2 = 0; q2 < q_head_per_kv_head_count; q2++) {
+     for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+      int h = h2;
+      for (int p = 0; p <= i + bi; p++) {
+       mha_score[h2][q2][i][p] = 0.f;
+       for (int e = 0; e < head_dim; e++) {
+        mha_score[h2][q2][i][p] +=
+          mha_q[h2][q2][i][e] * k_cache[l][h][p][e];
+       }
+       mha_score[h2][q2][i][p] /= sqrtf(head_dim);
+      }
+      softmax(i + bi + 1, context_len, &mha_score[h2][q2][i][0]);
+      for (int e = 0; e < head_dim; e++) {
+       mha_blend[h2][q2][i][e] = 0.f;
+      }
+      for (int p = 0; p <= i + bi; p++) {
+       for (int e = 0; e < head_dim; e++) {
+        mha_blend[h2][q2][i][e] +=
+          mha_score[h2][q2][i][p] * v_cache[l][h][p][e];
+       }
+      }
+     }
+    }
+   }
+   for (int h2 = 0; h2 < kv_head_count; h2++) {
+    for (int q2 = 0; q2 < q_head_per_kv_head_count; q2++) {
+     for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+      for (int e = 0; e < head_dim; e++) {
+       mha_att[i][(h2 * q_head_per_kv_head_count + q2) * head_dim + e] =
+         mha_blend[h2][q2][i][e];
+      }
+     }
+    }
+   }
+   matmul(min(sequence_len, bi + 512) - bi, embedding_dim, embedding_dim,
+       mha_out, mha_att, &mha_out_weight[l][0][0]);
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    for (int e = 0; e < embedding_dim; e++) {
+     embedding[i][e] += mha_out[i][e];
+    }
+   }
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    rmsnorm(embedding_dim, &ffn_norm[i][0], &embedding[i][0],
+        &ffn_norm_weight[l][0], epsilon);
+   }
+   matmul(min(sequence_len, bi + 512) - bi, hidden_dim, embedding_dim,
+       ffn_fc, ffn_norm, &ffn_fc_weight[l][0][0]);
+   matmul(min(sequence_len, bi + 512) - bi, hidden_dim, embedding_dim,
+       ffn_up, ffn_norm, &ffn_up_weight[l][0][0]);
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    for (int e = 0; e < hidden_dim; e++) {
+     ffn_fc[i][e] *= 1.f / (1.f + expf(-ffn_fc[i][e]));
+     ffn_fc[i][e] *= ffn_up[i][e];
+    }
+   }
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    for (int j = 0; j < embedding_dim; j++) {
+     ffn_out[i][j] = 0.f;
+     for (int k = 0; k < hidden_dim; k++) {
+      ffn_out[i][j] += ffn_fc[i][k] * ffn_out_weight[l][j][k];
+     }
+    }
+   }
+   for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+    for (int e = 0; e < embedding_dim; e++) {
+     embedding[i][e] += ffn_out[i][e];
+    }
+   }
+  }
+  for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+   rmsnorm(embedding_dim, &embedding[i][0], &embedding[i][0],
+       out_norm_weight, epsilon);
+  }
+  for (int i = 0; i < min(sequence_len, bi + 512) - bi; i++) {
+   if (i + bi == sequence_len - 1 ? 1 : 0) {
+    for (int j = 0; j < vocabulary_len; j++) {
+     logits[j] = 0.f;
+     for (int k = 0; k < embedding_dim; k++) {
+      logits[j] += embedding[i][k] * out_weight[j][k];
+     }
+    }
+   }
+  }
+  free(ffn_out);
+  free(ffn_up);
+  free(ffn_fc);
+  free(ffn_norm);
+  free(mha_out);
+  free(mha_att);
+  free(mha_blend);
+  free(mha_score);
+  free(mha_q);
+  free(mha_norm);
+  free(embedding);
  }
 }