[examples/SplitLlama] Adjust code structure, add build runtime related interfaces, reduce duplicate code

examples/BuddyLlama/llama-main.cpp

@@ -129,9 +129,7 @@ int main() {
   //  - Output container.
   //  - Parameters container.
   Text<size_t, 2> outputContainer;
-  MemRef<float, 3> resultContainer[2] = {
-      MemRef<float, 3>({1, MaxTokenLength, HiddenSize}, false, 0),
-      MemRef<float, 3>({1, MaxTokenLength, MaxVocabSize}, false, 0)};
+  MemRef<float, 3> resultContainer({1, MaxTokenLength, HiddenSize}, false, 0);
   Text<size_t, 2> inputContainer(inputStr);
   MemRef<float, 1> paramsContainer({ParamsSize});
 
@@ -151,7 +149,7 @@ int main() {
   for (int i = 0; i < generateLen; i++) {
     const auto inferenceStart = std::chrono::high_resolution_clock::now();
     // Execute the forward pass of the model.
-    _mlir_ciface_forward(resultContainer, &paramsContainer, &inputContainer);
+    _mlir_ciface_forward(&resultContainer, &paramsContainer, &inputContainer);
 
     const auto inferenceEnd = std::chrono::high_resolution_clock::now();
     const std::chrono::duration<double, std::milli> inferenceTime =
@@ -160,7 +158,7 @@ int main() {
     // Determine the generated token.
     int tokenIndex = inputContainer.getTokenCnt() - 1;
     const float *startPtr =
-        resultContainer[1].getData() + tokenIndex * MaxVocabSize;
+        resultContainer.getData() + tokenIndex * MaxVocabSize;
     const float *endPtr = startPtr + MaxVocabSize;
     int maxIndex = findMaxIndex(startPtr, endPtr);
     std::string tok = inputContainer.getStr(maxIndex);
@@ -175,8 +173,7 @@ int main() {
     // Append the generated token into the input and output container.
     inputContainer.appendTokenIdx(maxIndex);
     outputContainer.appendTokenIdx(maxIndex);
-    free(resultContainer[0].release());
-    free(resultContainer[1].release());
+    free(resultContainer.release());
   }
 
   /// Print the final result

examples/CMakeLists.txt

@@ -12,6 +12,10 @@ if(BUDDY_DIS_LLAMA_EXAMPLES)
   add_subdirectory(LlamaTest)
 endif()
 
+if(BUDDY_SPLIT_LLAMA_EXAMPLES)
+  add_subdirectory(SplitLlama)
+endif()
+
 if(BUDDY_DEEPSEEKR1_EXAMPLES)
   add_subdirectory(BuddyDeepSeekR1)
 endif()

examples/LlamaTest/.gitignore

@@ -1,5 +1,5 @@
 # model params file
-arg0.data
+*.data
 vocab.txt
 
 # model mlir file