From bb5ea10ecb24adb176b65da30780cf02b7ee58b7 Mon Sep 17 00:00:00 2001
From: Lin Jiang <linjiang@taichi.graphics>
Date: Tue, 20 Jun 2023 17:05:28 +0800
Subject: [PATCH] [llvm] Allocate the buffers of the real function in the entry
 block (#8206)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Issue: fix #8203

### Brief Summary

<!--
copilot:summary
-->
### <samp>🤖 Generated by Copilot at 4f313de</samp>

Refactor the LLVM backend to use a better memory allocation strategy
with `create_entry_block_alloca`. This improves performance and
correctness of the generated code.

### Walkthrough

<!--
copilot:walkthrough
-->
### <samp>🤖 Generated by Copilot at 4f313de</samp>

* Use `create_entry_block_alloca` function to allocate memory for LLVM
variables
([link](https://github.com/taichi-dev/taichi/pull/8206/files?diff=unified&w=0#diff-3c663c78745adcd3f6a7ac81fe99e628decc3040f292ea1e20ecd4b85a7f4313L2784-R2788),
[link](https://github.com/taichi-dev/taichi/pull/8206/files?diff=unified&w=0#diff-3c663c78745adcd3f6a7ac81fe99e628decc3040f292ea1e20ecd4b85a7f4313L2795-R2795)).
This avoids potential stack size and alignment issues in the LLVM
backend. The variables are `new_ctx`, `buffer`, and `result_buffer` in
`codegen_llvm.cpp`.
---
 taichi/codegen/llvm/codegen_llvm.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/taichi/codegen/llvm/codegen_llvm.cpp b/taichi/codegen/llvm/codegen_llvm.cpp
index 79c03d5d48b36..d520b8dc38559 100644
--- a/taichi/codegen/llvm/codegen_llvm.cpp
+++ b/taichi/codegen/llvm/codegen_llvm.cpp
@@ -2781,18 +2781,18 @@ void TaskCodeGenLLVM::visit(FuncCallStmt *stmt) {
     current_callable = old_callable;
   }
   llvm::Function *llvm_func = func_map[stmt->func];
-  auto *new_ctx = builder->CreateAlloca(get_runtime_type("RuntimeContext"));
+  auto *new_ctx = create_entry_block_alloca(get_runtime_type("RuntimeContext"));
   call("RuntimeContext_set_runtime", new_ctx, get_runtime());
   if (!stmt->func->parameter_list.empty()) {
     auto *buffer =
-        builder->CreateAlloca(tlctx->get_data_type(stmt->func->args_type));
+        create_entry_block_alloca(tlctx->get_data_type(stmt->func->args_type));
     set_args_ptr(stmt->func, new_ctx, buffer);
     set_struct_to_buffer(stmt->func->args_type, buffer, stmt->args);
   }
   llvm::Value *result_buffer = nullptr;
   if (!stmt->func->rets.empty()) {
     auto *ret_type = tlctx->get_data_type(stmt->func->ret_type);
-    result_buffer = builder->CreateAlloca(ret_type);
+    result_buffer = create_entry_block_alloca(ret_type);
     auto *result_buffer_u64 = builder->CreatePointerCast(
         result_buffer,
         llvm::PointerType::get(tlctx->get_data_type<uint64>(), 0));