/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file codegen_cpu.cc
 */
#ifdef TVM_LLVM_VERSION

#include "codegen_cpu.h"

#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringRef.h>
#include <llvm/IR/Argument.h>
#include <llvm/IR/Attributes.h>
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/CallingConv.h>
#include <llvm/IR/Comdat.h>
#include <llvm/IR/Constants.h>
#include <llvm/IR/DIBuilder.h>
#include <llvm/IR/DebugInfoMetadata.h>
#include <llvm/IR/DebugLoc.h>
#include <llvm/IR/DerivedTypes.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/GlobalVariable.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/MDBuilder.h>
#include <llvm/IR/Metadata.h>
#include <llvm/IR/Module.h>
#include <llvm/Support/Alignment.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Target/TargetMachine.h>
#include <llvm/Transforms/Utils/ModuleUtils.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/analysis.h>

#include <algorithm>
#include <memory>
#include <tuple>
#include <unordered_map>
#include <unordered_set>

#include "llvm_instance.h"

namespace tvm {
namespace codegen {

// Make these non-inline because of std::unique_ptr. See comment in
// codegen_llvm.cc for more information.
CodeGenCPU::CodeGenCPU() = default;
CodeGenCPU::~CodeGenCPU() = default;

void CodeGenCPU::Init(const std::string& module_name, LLVMTarget* llvm_target,
                      ffi::Optional<ffi::String> system_lib_prefix, bool dynamic_lookup,
                      bool target_c_runtime) {
  CodeGenLLVM::Init(module_name, llvm_target, system_lib_prefix, dynamic_lookup, target_c_runtime);
  system_lib_prefix_ = system_lib_prefix;
  dbg_info_ = CreateDebugInfo(module_.get());
  func_handle_map_.clear();
  func_handle_init_map_.clear();
  export_system_symbols_.clear();

  // Runtime types.
  t_tvm_shape_index_ =
      llvm::Type::getIntNTy(*llvm_target_->GetContext(), DefaultIndexPrimType().bits());
  // Defined in 3rdparty/dlpack/include/dlpack/dlpack.h:
  // typedef struct { DLDeviceType device_type; int device_id; } DLDevice;
  t_tvm_device_ = llvm::StructType::create({t_int_, t_int_});
  // Defined in 3rdparty/dlpack/include/dlpack/dlpack.h:
  // typedef struct { uint8_t code; uint8_t bits; uint16_t lanes; } DLDataType;
  t_tvm_type_ = llvm::StructType::create({t_int8_, t_int8_, t_int16_});
  // Defined in include/tvm/runtime/base.h:
  // typedef void* TVMFunctionHandle;
  t_tvm_func_handle_ = t_void_p_;
  // Defined in 3rdparty/dlpack/include/dlpack/dlpack.h:
  // typedef struct { ... } DLTensor;
  t_tvm_array_ = llvm::StructType::create({t_void_p_, t_tvm_device_, t_int_, t_tvm_type_,
                                           llvmGetPointerTo(t_tvm_shape_index_, 0),
                                           llvmGetPointerTo(t_tvm_shape_index_, 0), t_int64_});
  // Defined in include/tvm/ffi/c_api.h:
  t_tvm_ffi_any_ = llvm::StructType::create({t_int32_, t_int32_, t_float64_});
  // Defined in include/tvm/runtime/c_backend_api.h:
  // typedef struct { void* sync_handle; int32_t num_task; } TVMParallelGroupEnv;
  t_tvm_parallel_group_env_ = llvm::StructType::create({llvmGetPointerTo(t_int32_, 0), t_int32_});
  // Defined in include/tvm/ffi/c_api.h:
  // typedef int (*)(void* self, const TVMFFIAny* args, int32_t num_args,
  //                 TVMFFIAny* result);
  ftype_tvm_ffi_c_func_ = llvm::FunctionType::get(
      t_int_,
      {t_void_p_, llvmGetPointerTo(t_tvm_ffi_any_, 0), t_int_, llvmGetPointerTo(t_tvm_ffi_any_, 0)},
      false);
  // Defined in include/tvm/runtime/c_backend_api.h:
  // typedef int (*FTVMParallelLambda)(int task_id, TVMParallelGroupEnv* penv, void* cdata);
  ftype_tvm_parallel_lambda_ = llvm::FunctionType::get(
      t_int_, {t_int_, llvmGetPointerTo(t_tvm_parallel_group_env_, 0), t_void_p_}, false);
  md_tbaa_ctx_ptr_ = md_builder_->createTBAAScalarTypeNode("ctx_ptr", md_tbaa_root_);

  // Runtime functions.
  // Defined in include/tvm/ffi/c_api.h:
  // int TVMFFIFunctionCall(TVMFunctionHandle func, TVMFFIAny* args, int32_t num_args,
  //                    TVMFFIAny* result);
  ftype_tvm_ffi_func_call_ = ftype_tvm_ffi_c_func_;
  // Defined in include/tvm/ffi/c_api.h:
  // void TVMFFIErrorSetRaisedFromCStr(const char *kind, const char* msg);
  ftype_tvm_ffi_error_set_raised_by_c_str_ = llvm::FunctionType::get(
      t_void_, {llvmGetPointerTo(t_char_, 0), llvmGetPointerTo(t_char_, 0)}, false);
  // Defined in include/tvm/ffi/c_api.h:
  // void TVMFFIErrorSetRaisedFromCStrParts(const char* kind, const char** msg_parts, int32_t n);
  ftype_tvm_ffi_error_set_raised_from_c_str_parts_ = llvm::FunctionType::get(
      t_void_,
      {llvmGetPointerTo(t_char_, 0), llvmGetPointerTo(llvmGetPointerTo(t_char_, 0), 0), t_int_},
      false);
  // Defined in include/tvm/ffi/extra/c_env_api.h:
  // int TVMFFIEnvModLookupFromImports(TVMFFIObjectHandle library_ctx, const char* func_name,
  //                                   TVMFFIObjectHandle* out);
  ftype_tvm_ffi_env_mod_lookup_from_imports_ = llvm::FunctionType::get(
      t_int_, {t_void_p_, llvmGetPointerTo(t_char_, 0), llvmGetPointerTo(t_tvm_func_handle_, 0)},
      false);
  // Defined in include/tvm/ffi/c_api.h:
  // int TVMFFIHandleInitOnce(void** handle_addr, int (*init_func)(void** result));
  ftype_tvm_ffi_handle_init_callback_ =
      llvm::FunctionType::get(t_int_, {llvmGetPointerTo(t_void_p_, 0)}, false);
  ftype_tvm_ffi_handle_init_once_ = llvm::FunctionType::get(
      t_int_,
      {llvmGetPointerTo(t_void_p_, 0), llvmGetPointerTo(ftype_tvm_ffi_handle_init_callback_, 0)},
      false);
  // Defined in include/tvm/runtime/c_backend_api.h:
  // int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_task);
  ftype_tvm_parallel_launch_ = llvm::FunctionType::get(
      t_int_, {llvmGetPointerTo(ftype_tvm_parallel_lambda_, 0), t_void_p_, t_int_}, false);
  // Defined in include/tvm/runtime/c_backend_api.h:
  // int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv);
  ftype_tvm_parallel_barrier_ = llvm::FunctionType::get(
      t_int_, {t_int_, llvmGetPointerTo(t_tvm_parallel_group_env_, 0)}, false);
  // initialize TVM runtime API
  if (system_lib_prefix_.has_value() && !target_c_runtime) {
    // We will need this in environment for backward registration.
    // Defined in include/tvm/runtime/c_backend_api.h:
    // int TVMFFIEnvModRegisterSystemLibSymbol(const char* name, void* ptr);
    f_tvm_register_system_symbol_ = llvm::Function::Create(
        llvm::FunctionType::get(t_int_, {llvmGetPointerTo(t_char_, 0), t_void_p_}, false),
        llvm::Function::ExternalLinkage, "TVMFFIEnvModRegisterSystemLibSymbol", module_.get());
  } else {
    f_tvm_register_system_symbol_ = nullptr;
  }
  if (dynamic_lookup || system_lib_prefix_.has_value()) {
    f_tvm_ffi_func_call_ =
        llvm::Function::Create(ftype_tvm_ffi_func_call_, llvm::Function::ExternalLinkage,
                               "TVMFFIFunctionCall", module_.get());
    f_tvm_ffi_set_raised_by_c_str_ = llvm::Function::Create(
        ftype_tvm_ffi_error_set_raised_by_c_str_, llvm::Function::ExternalLinkage,
        "TVMFFIErrorSetRaisedFromCStr", module_.get());
    f_tvm_ffi_set_raised_from_c_str_parts_ = llvm::Function::Create(
        ftype_tvm_ffi_error_set_raised_from_c_str_parts_, llvm::Function::ExternalLinkage,
        "TVMFFIErrorSetRaisedFromCStrParts", module_.get());
    f_tvm_ffi_env_mod_lookup_from_imports_ = llvm::Function::Create(
        ftype_tvm_ffi_env_mod_lookup_from_imports_, llvm::Function::ExternalLinkage,
        "TVMFFIEnvModLookupFromImports", module_.get());
    f_tvm_ffi_handle_init_once_ =
        llvm::Function::Create(ftype_tvm_ffi_handle_init_once_, llvm::Function::ExternalLinkage,
                               "TVMFFIHandleInitOnce", module_.get());
    f_tvm_parallel_launch_ =
        llvm::Function::Create(ftype_tvm_parallel_launch_, llvm::Function::ExternalLinkage,
                               "TVMBackendParallelLaunch", module_.get());
    f_tvm_parallel_barrier_ =
        llvm::Function::Create(ftype_tvm_parallel_barrier_, llvm::Function::ExternalLinkage,
                               "TVMBackendParallelBarrier", module_.get());
  }
  target_c_runtime_ = target_c_runtime;
  InitGlobalContext(dynamic_lookup);
}

llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(llvm::StringRef name,
                                                    const ffi::Array<Type>& param_types,
                                                    const Type& return_type) {
  llvm::SmallVector<llvm::Metadata*, 4> paramTys;

  paramTys.push_back(GetDebugType(return_type));
  for (const auto& param_type : param_types) {
    paramTys.push_back(GetDebugType(param_type));
  }

  auto* DIFunctionTy = dbg_info_->di_builder_->createSubroutineType(
      dbg_info_->di_builder_->getOrCreateTypeArray(paramTys));

  bool local_to_unit = llvm::GlobalVariable::isLocalLinkage(llvm::GlobalValue::InternalLinkage);

  auto SPFlags = llvm::DISubprogram::toSPFlags(local_to_unit, /*IsDefinition=*/true,
                                               /*IsOptimized=*/true);

  auto* DIFunction = dbg_info_->di_builder_->createFunction(
      /*Scope=*/dbg_info_->file_, /*Name=*/name, /*LinkageName=*/"",
      /*File=*/dbg_info_->file_, /*LineNo=*/0, /*Ty=*/DIFunctionTy,
      /*ScopeLine=*/0, /*Flags=*/llvm::DINode::FlagPrototyped, /*SPFlags=*/SPFlags);

  return DIFunction;
}

llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const GlobalVar& gvar, const PrimFunc& func) {
  std::string name = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint);
  return CreateDebugFunction(name, func->params.Map([](const Var& var) { return var->ty; }),
                             func->ret_type);
}

void CodeGenCPU::AddFunction(const GlobalVar& gvar, const PrimFunc& func) {
  di_subprogram_ = CreateDebugFunction(gvar, func);
  EmitDebugLocation(func->span);
  CodeGenLLVM::AddFunction(gvar, func);
  if (f_tvm_register_system_symbol_ != nullptr) {
    if (auto global_symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol)) {
      export_system_symbols_.emplace_back(
          std::make_pair(global_symbol.value().operator std::string(), function_));
    }
  }
  AddDebugInformation(function_, func->params.Map([](const Var& var) { return var->ty; }));
}

void CodeGenCPU::AddMainFunction(const std::string& entry_func_name) {
  if (module_->getFunction(ffi::symbol::tvm_ffi_main) != nullptr) {
    // main already exists, no need to create a wrapper function
    // main takes precedence over other entry functions
    return;
  }
  // create a wrapper function with tvm_ffi_main name and redirects to the entry function
  llvm::Function* target_func = module_->getFunction(entry_func_name);
  TVM_FFI_ICHECK(target_func) << "Function " << entry_func_name << " does not exist in module";

  // Create wrapper function
  llvm::Function* wrapper_func =
      llvm::Function::Create(target_func->getFunctionType(), llvm::Function::WeakAnyLinkage,
                             ffi::symbol::tvm_ffi_main, module_.get());

  // Set attributes (Windows comdat, DLL export, etc.)
  if (llvm_target_->GetOrCreateTargetMachine()->getTargetTriple().isOSWindows()) {
    llvm::Comdat* comdat = module_->getOrInsertComdat(ffi::symbol::tvm_ffi_main);
    comdat->setSelectionKind(llvm::Comdat::Any);
    wrapper_func->setComdat(comdat);
  }

  wrapper_func->setCallingConv(llvm::CallingConv::C);
  wrapper_func->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);

  // Create simple tail call
  llvm::BasicBlock* entry =
      llvm::BasicBlock::Create(*llvm_target_->GetContext(), "entry", wrapper_func);
  builder_->SetInsertPoint(entry);

  // Forward all arguments to target function
  std::vector<llvm::Value*> call_args;
  for (llvm::Value& arg : wrapper_func->args()) {
    call_args.push_back(&arg);
  }

  llvm::Value* result = builder_->CreateCall(target_func, call_args);
  if (target_func->getReturnType()->isVoidTy()) {
    builder_->CreateRetVoid();
  } else {
    builder_->CreateRet(result);
  }
}

std::unique_ptr<llvm::Module> CodeGenCPU::Finish() {
  // link modules
  if (dbg_info_ != nullptr) {
    dbg_info_->di_builder_->finalize();
  }
  return CodeGenLLVM::Finish();
}

CodeGenLLVM::TypedPointer CodeGenCPU::CreateStructRefPtr(Type type, llvm::Value* buf,
                                                         llvm::Value* index, int kind) {
  if (kind < builtin::kDLTensorKindBound_) {
    if (buf->getType() == t_void_p_) {
      buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_array_, 0));
    } else {
      TVM_FFI_ICHECK_EQ(buf->getType(), llvmGetPointerTo(t_tvm_array_, 0));
    }
  }
  switch (kind) {
    case builtin::kDLTensorAddr: {
      return TypedPointer(t_tvm_array_, builder_->CreateInBoundsGEP(t_tvm_array_, buf, index));
    }
    case builtin::kDLTensorData: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(0);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(0)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorShape: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(4);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(4)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorStrides: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(5);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(5)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorNDim: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(2);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(2)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorTypeCode: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(0);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(0)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorTypeBits: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(1);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(1)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorTypeLanes: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(2);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(2)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorByteOffset: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(6);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(6)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorDeviceId: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(1)->getStructElementType(1);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(1), ConstInt32(1)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kDLTensorDeviceType: {
      llvm::Type* member_type = t_tvm_array_->getStructElementType(1)->getStructElementType(0);
      llvm::Value* member_addr =
          builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(1), ConstInt32(0)});
      return TypedPointer(member_type, member_addr);
    }
    case builtin::kTVMFFIAnyTypeIndex: {
      buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_ffi_any_, 0));
      buf = builder_->CreateInBoundsGEP(t_tvm_ffi_any_, buf, {index, ConstInt32(0)});
      return TypedPointer(t_int32_, buf);
    }
    case builtin::kTVMFFIAnyZeroPadding: {
      buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_ffi_any_, 0));
      buf = builder_->CreateInBoundsGEP(t_tvm_ffi_any_, buf, {index, ConstInt32(1)});
      return TypedPointer(t_int32_, buf);
    }
    case builtin::kTVMFFIAnyUnionValue: {
      buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_ffi_any_, 0));
      // field 2 is the union value
      buf = builder_->CreateInBoundsGEP(t_tvm_ffi_any_, buf, {index, ConstInt32(2)});
      if (type.as<PointerTypeNode>()) {
        buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_void_p_, 0));
        return TypedPointer(t_void_p_, buf);
      }

      PrimType dtype = type.as_or_throw<PrimType>();
      TVM_FFI_ICHECK_EQ(dtype.lanes(), 1);
      if (dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
        // it should be safe to set the pointer to the first byte of the union value
        buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(DTypeToLLVMType(dtype), 0));
        return TypedPointer(t_int8_, buf);
      } else if (dtype.MatchesCode(DLDataTypeCode::kDLInt) && dtype.bits() == 64) {
        buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_int64_, 0));
        return TypedPointer(t_int64_, buf);
      } else if (dtype.MatchesCode(DLDataTypeCode::kDLFloat) && dtype.bits() == 64) {
        buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_float64_, 0));
        return TypedPointer(t_float64_, buf);
      }
      TVM_FFI_THROW(InternalError)
          << "Type " << type << " cannot be stored into a TVMFFIAny's value field";
    }
    case builtin::kInt64ArrayElem: {
      buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_int64_, 0));
      llvm::Value* elem_addr = builder_->CreateInBoundsGEP(t_int64_, buf, index);
      return TypedPointer(t_int64_, elem_addr);
    }
    default:
      TVM_FFI_THROW(InternalError) << "unknown field code";
  }
}

llvm::Value* CodeGenCPU::CreateCallExtern(Type ret_type, ffi::String global_symbol,
                                          const ffi::Array<Expr>& args, bool skip_first_arg) {
  std::vector<llvm::Value*> arg_values;
  for (size_t i = static_cast<size_t>(skip_first_arg); i < args.size(); ++i) {
    arg_values.push_back(MakeValue(args[i]));
  }
  std::vector<llvm::Type*> arg_types;
  for (llvm::Value* v : arg_values) {
    arg_types.push_back(v->getType());
  }
  llvm::FunctionType* ftype = llvm::FunctionType::get(GetLLVMType(ret_type), arg_types, false);
  // Check if it is available in global function table as injected function.

  auto callee = [&]() -> llvm::Value* {
    if (auto it = gv_func_map_.find(global_symbol); it != gv_func_map_.end()) {
      if (it->second == nullptr) {
        it->second = InitContextPtr(llvmGetPointerTo(ftype, 0), "__" + global_symbol);
      }
      return GetContextPtr(it->second);
    } else if (llvm::Function* f = module_->getFunction(MakeStringRef(global_symbol))) {
      return f;
    } else {
      return llvm::Function::Create(ftype, llvm::Function::ExternalLinkage,
                                    MakeStringRef(global_symbol), module_.get());
    }
  }();

  if (callee->getType() != llvmGetPointerTo(ftype, 0)) {
    callee = builder_->CreatePointerCast(callee, llvmGetPointerTo(ftype, 0));
  }
  return builder_->CreateCall(ftype, callee, arg_values);
}

llvm::GlobalVariable* CodeGenCPU::InitContextPtr(llvm::Type* p_type, std::string name) {
  llvm::GlobalVariable* gv = new llvm::GlobalVariable(
      *module_, p_type, false, llvm::GlobalValue::LinkOnceAnyLinkage, nullptr, name);
  gv->setAlignment(llvm::Align(data_layout_->getTypeAllocSize(p_type)));
  gv->setInitializer(llvm::Constant::getNullValue(p_type));
  gv->setDLLStorageClass(llvm::GlobalValue::DLLStorageClassTypes::DLLExportStorageClass);
  // comdat is needed for windows select any linking to work
  // set comdat to Any(weak linking)
  if (llvm_target_->GetOrCreateTargetMachine()->getTargetTriple().isOSWindows()) {
    llvm::Comdat* comdat = module_->getOrInsertComdat(name);
    comdat->setSelectionKind(llvm::Comdat::Any);
    gv->setComdat(comdat);
  }
  return gv;
}

llvm::Value* CodeGenCPU::GetContextPtr(llvm::GlobalVariable* gv) {
  TVM_FFI_ICHECK(gv != nullptr);
  llvm::LoadInst* faddr =
      builder_->CreateAlignedLoad(gv->getValueType(), gv, llvm::Align(gv->getAlignment()));
  faddr->setMetadata("tbaa",
                     md_builder_->createTBAAStructTagNode(md_tbaa_ctx_ptr_, md_tbaa_ctx_ptr_, 0));
  return faddr;
}

void CodeGenCPU::InitGlobalContext(bool dynamic_lookup) {
  std::string ctx_symbol = system_lib_prefix_.value_or("") + ffi::symbol::tvm_ffi_library_ctx;
  // Module context
  gv_mod_ctx_ = InitContextPtr(t_void_p_, ctx_symbol);
  // Register back the locations.
  if (f_tvm_register_system_symbol_ != nullptr && !target_c_runtime_) {
    export_system_symbols_.emplace_back(std::make_pair(ctx_symbol, gv_mod_ctx_));
  } else {
    if (!dynamic_lookup) {
      gv_tvm_ffi_func_call_ =
          InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_func_call_, 0), "__TVMFFIFunctionCall");
      gv_tvm_ffi_env_mod_lookup_from_imports_ =
          InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_env_mod_lookup_from_imports_, 0),
                         "__TVMFFIEnvModLookupFromImports");
      gv_tvm_ffi_handle_init_once_ = InitContextPtr(
          llvmGetPointerTo(ftype_tvm_ffi_handle_init_once_, 0), "__TVMFFIHandleInitOnce");
      gv_tvm_ffi_set_last_error_c_str_ =
          InitContextPtr(llvmGetPointerTo(ftype_tvm_ffi_error_set_raised_by_c_str_, 0),
                         "__TVMFFIErrorSetRaisedFromCStr");
      gv_tvm_parallel_launch_ = InitContextPtr(llvmGetPointerTo(ftype_tvm_parallel_launch_, 0),
                                               "__TVMBackendParallelLaunch");
      gv_tvm_parallel_barrier_ = InitContextPtr(llvmGetPointerTo(ftype_tvm_parallel_barrier_, 0),
                                                "__TVMBackendParallelBarrier");
      // Mark as context functions
      gv_func_map_["TVMBackendAllocWorkspace"] = nullptr;
      gv_func_map_["TVMBackendFreeWorkspace"] = nullptr;
    }
  }
}

llvm::BasicBlock* CodeGenCPU::CheckCallSuccess(llvm::Value* retcode) {
  // create emit codes that checks and load the function.
  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  auto* fail_block = llvm::BasicBlock::Create(*ctx, "call_fail", function_);
  auto* end_block = llvm::BasicBlock::Create(*ctx, "call_end", function_);
  auto* succ = builder_->CreateICmpEQ(retcode, llvm::ConstantInt::get(t_int_, 0));
  builder_->CreateCondBr(succ, end_block, fail_block, md_very_likely_branch_);
  builder_->SetInsertPoint(fail_block);
  // return the code.
  builder_->CreateRet(retcode);
  // otherwise set it to be new end.
  builder_->SetInsertPoint(end_block);
  return end_block;
}

void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) {
  EmitDebugLocation(op);
  /*! \brief maintain states that should be guarded when step into compute scope */
  struct ComputeScopeStates {
    explicit ComputeScopeStates(CodeGenCPU* parent) : parent_(parent) {}

    void EnterWithScope() {
      std::swap(function_, parent_->function_);
      std::swap(analyzer_, parent_->analyzer_);
      std::swap(var_map_, parent_->var_map_);
      std::swap(di_subprogram_, parent_->di_subprogram_);
      std::swap(loop_frame_jump_tgts_, parent_->loop_frame_jump_tgts_);
    }

    void ExitWithScope() {
      std::swap(function_, parent_->function_);
      std::swap(analyzer_, parent_->analyzer_);
      std::swap(var_map_, parent_->var_map_);
      std::swap(di_subprogram_, parent_->di_subprogram_);
      std::swap(loop_frame_jump_tgts_, parent_->loop_frame_jump_tgts_);
    }

    llvm::Function* function_{nullptr};
    llvm::DISubprogram* di_subprogram_{nullptr};
    std::unordered_map<const VarNode*, llvm::Value*> var_map_;
    std::vector<std::pair<llvm::BasicBlock*, llvm::BasicBlock*>> loop_frame_jump_tgts_;
    arith::Analyzer analyzer_{arith::Analyzer()};
    CodeGenCPU* parent_;
  };

  // There are two reasons why we create another function for compute_scope
  // - Make sure the generated compute function is clearly separately(though it can get inlined)
  // - Set noalias on all the pointer arguments, some of them are loaded from ffi::PackedArgs.
  //   This is easier than set the alias scope manually.
  ffi::Array<Var> vargs = tirx::UndefinedVars(op->body, {});
  std::vector<llvm::Value*> arg_values;
  std::vector<llvm::Type*> arg_types;
  for (Var v : vargs) {
    llvm::Value* value = MakeValue(v);
    value->setName(v->name_hint.c_str());
    arg_values.push_back(value);
    arg_types.push_back(value->getType());
  }
  llvm::FunctionType* ftype = llvm::FunctionType::get(t_int_, arg_types, false);
  // $xxx_compute_ functions are not global. They should be marked as static (via InternalLinkage)
  // to call them correctly on MIPS platform (CALL16 issue)
  // Linkage ld Error: CALL16 reloc at 0x290 not against global symbol
  const StringImmNode* value = op->value.as<StringImmNode>();
  TVM_FFI_ICHECK(value != nullptr);
  llvm::Function* fcompute = llvm::Function::Create(ftype, llvm::Function::InternalLinkage,
                                                    MakeStringRef(value->value), module_.get());
  SetTargetAttributes(fcompute);
  for (auto it = fcompute->arg_begin(); it != fcompute->arg_end(); it++) {
    const Var& var = vargs[std::distance(fcompute->arg_begin(), it)];
    it->setName(std::string(var->name_hint));
  }

  llvm::BasicBlock* compute_call_end = CheckCallSuccess(builder_->CreateCall(fcompute, arg_values));
  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  // enter compute scope and setup compute function.
  With<ComputeScopeStates> scope_states_guard(this);
  size_t idx = 0;
  for (auto it = fcompute->arg_begin(); it != fcompute->arg_end(); ++it, ++idx) {
    llvm::Argument* v = &(*it);
    const Var& var = vargs[idx];
    var_map_[var.get()] = v;
    if (var->ty.as<PointerTypeNode>() && !alias_var_set_.count(var.get())) {
      // set non alias.
      fcompute->addParamAttr(idx, llvm::Attribute::NoAlias);
      // always not inline compute function to make the code structure clean
      fcompute->addFnAttr(llvm::Attribute::NoInline);
    }
    // Add alignment attribute if needed.
    auto f = alloc_storage_info_.find(var.get());
    if (f != alloc_storage_info_.end()) {
      unsigned align = f->second.alignment;
      if (align > 1) {
        auto attr = llvm::Attribute::get(*ctx, llvm::Attribute::Alignment, align);
        fcompute->addParamAttr(idx, attr);
      }
    }
  }

  function_ = fcompute;
  di_subprogram_ =
      CreateDebugFunction(MakeStringRef(value->value),
                          vargs.Map([](const Var& var) { return var->ty; }), PrimType::Int(32));
  auto* compute_entry = llvm::BasicBlock::Create(*ctx, "entry", function_);
  builder_->SetInsertPoint(compute_entry);
  this->VisitStmt(op->body);
  builder_->CreateRet(ConstInt32(0));
  builder_->SetInsertPoint(compute_call_end);

  AddDebugInformation(fcompute, vargs.Map([](const Var& var) { return var->ty; }));
}

CodeGenLLVM::TypedPointer CodeGenCPU::PackClosureData(const ffi::Array<Var>& vfields,
                                                      uint64_t* num_bytes,
                                                      std::string struct_name) {
  if (vfields.size() == 0) {
    *num_bytes = 0U;
    return TypedPointer(t_void_p_, llvm::Constant::getNullValue(t_void_p_));
  }
  std::vector<llvm::Type*> fields;
  for (Var v : vfields) {
    auto it = var_map_.find(v.get());
    TVM_FFI_ICHECK(it != var_map_.end());
    fields.push_back(it->second->getType());
  }
  llvm::StructType* ctype = struct_name.size() ? llvm::StructType::create(fields, struct_name)
                                               : llvm::StructType::create(fields);
  llvm::AllocaInst* cvalue =
      WithFunctionEntry([&]() { return builder_->CreateAlloca(ctype, ConstInt32(1)); });
  llvm::Value* zero = ConstInt32(0);
  for (size_t i = 0; i < vfields.size(); ++i) {
    builder_->CreateStore(var_map_.at(vfields[i].get()),
                          builder_->CreateInBoundsGEP(ctype, cvalue, {zero, ConstInt32(i)}));
  }
  *num_bytes = data_layout_->getTypeAllocSize(ctype);
  return TypedPointer(ctype, cvalue);
}

void CodeGenCPU::UnpackClosureData(TypedPointer cdata, const ffi::Array<Var>& vfields,
                                   std::unordered_map<const VarNode*, llvm::Value*>* vmap) {
  for (size_t i = 0; i < vfields.size(); ++i) {
    llvm::Type* field_type = cdata.type->getStructElementType(i);
    llvm::Value* field_addr =
        builder_->CreateInBoundsGEP(cdata.type, cdata.addr, {ConstInt32(0), ConstInt32(i)});
    llvm::Value* load =
        builder_->CreateLoad(field_type, field_addr, std::string(vfields[i]->name_hint));
    (*vmap)[vfields[i].get()] = load;
  }
}

void CodeGenCPU::CreateParallelLaunch(const Stmt& body, int num_task, std::string name) {
  // closure data
  llvm::Function* f =
      llvm::Function::Create(ftype_tvm_parallel_lambda_, llvm::Function::PrivateLinkage,
                             "__tvm_parallel_lambda", module_.get());
  SetTargetAttributes(f);

  // allocate and setup the closure, call the closure.
  ffi::Array<Var> vfields = tirx::UndefinedVars(body, {});
  uint64_t nbytes;
  TypedPointer cdata = PackClosureData(vfields, &nbytes, "closure_" + name);
  auto launch_callee = llvm::FunctionCallee(ftype_tvm_parallel_launch_, RuntimeTVMParallelLaunch());
  llvm::BasicBlock* par_launch_end = CheckCallSuccess(builder_->CreateCall(
      launch_callee,
      {f, builder_->CreatePointerCast(cdata.addr, t_void_p_), ConstInt32(num_task)}));
  // Setup the closure function.
  auto* lambda_entry =
      llvm::BasicBlock::Create(*llvm_target_->GetContext(), "parallel_closure_entry", f);
  builder_->SetInsertPoint(lambda_entry);
  auto it = f->arg_begin();
  llvm::Value* task_id = &(*it++);
  task_id->setName("task_id");
  llvm::Value* penv = &(*it++);
  cdata.addr = builder_->CreatePointerCast(&(*it++), cdata.addr->getType());
  // setup new variable map, swap it with current var context.
  std::unordered_map<const VarNode*, llvm::Value*> new_vmap;
  UnpackClosureData(cdata, vfields, &new_vmap);
  // setup parallel env
  ParallelEnv par_env;
  par_env.task_id = Var("task_id", PrimType::Int(32));
  par_env.num_task = Var("num_task", PrimType::Int(32));
  new_vmap[par_env.task_id.get()] = task_id;
  new_vmap[par_env.num_task.get()] = builder_->CreateLoad(
      t_int32_,
      builder_->CreateInBoundsGEP(t_tvm_parallel_group_env_, penv, {ConstInt32(0), ConstInt32(1)}),
      "num_task");
  par_env.penv = penv;
  auto new_analyzer = arith::Analyzer();
  std::swap(function_, f);
  std::swap(parallel_env_, par_env);
  std::swap(analyzer_, new_analyzer);
  std::swap(var_map_, new_vmap);
  this->VisitStmt(body);
  builder_->CreateRet(ConstInt32(0));
  // swap the var map back, now we are back on track.
  std::swap(var_map_, new_vmap);
  std::swap(analyzer_, new_analyzer);
  std::swap(parallel_env_, par_env);
  std::swap(function_, f);
  TVM_FFI_ICHECK_NE(par_env.parallel_loop_count, 0)
      << "Cannot find parallel loop within parallel launch";
  builder_->SetInsertPoint(par_launch_end);
}

llvm::Value* CodeGenCPU::CreateStaticHandle() {
  llvm::GlobalVariable* gv =
      new llvm::GlobalVariable(*module_, t_void_p_, false, llvm::GlobalValue::PrivateLinkage,
                               nullptr, "__tvm_static_handle");
  gv->setAlignment(llvm::Align(data_layout_->getTypeAllocSize(t_void_p_)));
  gv->setInitializer(llvm::Constant::getNullValue(t_void_p_));
  return gv;
}

llvm::Function* CodeGenCPU::CreatePackedFuncInit(const std::string& fname) {
  // Generate the equivalent of:
  //
  //   static int __tvm_func_handle_init.<fname>(void** result) {
  //     return TVMFFIEnvModLookupFromImports(__tvm_ffi__library_ctx, "<fname>", result);
  //   }
  //
  // The lookup result is owned by the module and remains valid for the generated module's lifetime.
  llvm::IRBuilderBase::InsertPoint saved_ip = builder_->saveIP();
  llvm::Function* saved_function = function_;
  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  llvm::Function* init_func =
      llvm::Function::Create(ftype_tvm_ffi_handle_init_callback_, llvm::Function::InternalLinkage,
                             "__tvm_func_handle_init." + fname, module_.get());
  init_func->setVisibility(llvm::GlobalValue::HiddenVisibility);
  SetTargetAttributes(init_func);
  function_ = init_func;
  builder_->SetInsertPoint(llvm::BasicBlock::Create(*ctx, "entry", init_func));

  llvm::Value* result = &*init_func->arg_begin();
  llvm::Value* module_ctx = builder_->CreateAlignedLoad(gv_mod_ctx_->getValueType(), gv_mod_ctx_,
                                                        llvm::Align(gv_mod_ctx_->getAlignment()));
  auto lookup_callee = llvm::FunctionCallee(ftype_tvm_ffi_env_mod_lookup_from_imports_,
                                            RuntimeTVMFFIEnvModLookupFromImports());
  builder_->CreateRet(
      builder_->CreateCall(lookup_callee, {module_ctx, GetConstString(fname), result}));

  function_ = saved_function;
  builder_->restoreIP(saved_ip);
  return init_func;
}

llvm::Value* CodeGenCPU::GetPackedFuncHandle(const std::string& fname) {
  // Cache the module-owned packed-function handle in global space with thread-safe initialization.
#if TVM_LLVM_VERSION >= 200
  llvm::DataLayout layout(module_.get()->getDataLayout());
#else
  llvm::DataLayout layout(module_.get());
#endif
  uint64_t align = layout.getTypeAllocSize(t_tvm_func_handle_);
  auto it = func_handle_map_.find(fname);

  llvm::GlobalVariable* hptr;
  llvm::Function* init_func;
  if (it == func_handle_map_.end()) {
    hptr =
        new llvm::GlobalVariable(*module_, t_tvm_func_handle_, false,
                                 llvm::GlobalValue::InternalLinkage, nullptr, ".tvm_func." + fname);
    hptr->setAlignment(llvm::Align(align));
    hptr->setInitializer(llvm::Constant::getNullValue(t_tvm_func_handle_));
    func_handle_map_[fname] = hptr;
    init_func = CreatePackedFuncInit(fname);
    func_handle_init_map_[fname] = init_func;
  } else {
    hptr = it->second;
    init_func = func_handle_init_map_.at(fname);
  }

  auto init_once_callee =
      llvm::FunctionCallee(ftype_tvm_ffi_handle_init_once_, RuntimeTVMFFIHandleInitOnce());
  CheckCallSuccess(builder_->CreateCall(init_once_callee, {hptr, init_func}));
  return builder_->CreateAlignedLoad(hptr->getValueType(), hptr, llvm::Align(align));
}

CodeGenCPU::PackedCall CodeGenCPU::MakeCallPackedLowered(const ffi::Array<Expr>& args,
                                                         const Type& r_type, const int64_t begin,
                                                         const int64_t end, bool use_env_lookup) {
  std::string func_name = [&]() {
    auto ptr = args[0].as<StringImmNode>();
    TVM_FFI_ICHECK(ptr) << "Expected first argument of Call to be "
                        << "a string containing the callee's name, "
                        << "but instead contained " << args[0];
    return ptr->value;
  }();
  // call the function
  int64_t nargs = end - begin;
  TVM_FFI_ICHECK_GE(nargs, 0);
  llvm::Value* stack_args = MakeValue(args[1]);
  llvm::Value* packed_args = builder_->CreateInBoundsGEP(
      t_tvm_ffi_any_, builder_->CreatePointerCast(stack_args, llvmGetPointerTo(t_tvm_ffi_any_, 0)),
      ConstInt32(begin));
  llvm::Value* result = builder_->CreateInBoundsGEP(
      t_tvm_ffi_any_, builder_->CreatePointerCast(stack_args, llvmGetPointerTo(t_tvm_ffi_any_, 0)),
      ConstInt32(end));

  llvm::FunctionType* callee_ftype = nullptr;
  llvm::Value* callee_value = nullptr;
  std::vector<llvm::Value*> call_args;

  if (use_env_lookup) {
    callee_ftype = ftype_tvm_ffi_func_call_;
    callee_value = RuntimeTVMFFIFunctionCall();
    call_args.push_back(GetPackedFuncHandle(func_name));
    call_args.insert(call_args.end(), {packed_args, ConstInt32(nargs), result});
  } else {
    // directly call into symbol, needs to prefix with tvm_ffi_symbol_prefix
    callee_ftype = ftype_tvm_ffi_c_func_;
    callee_value = module_->getFunction(ffi::symbol::tvm_ffi_symbol_prefix + func_name);
    if (callee_value == nullptr) {
      callee_value = llvm::Function::Create(ftype_tvm_ffi_c_func_, llvm::Function::ExternalLinkage,
                                            func_name, module_.get());
    }
    call_args.push_back(llvm::ConstantPointerNull::get(t_void_p_));
    call_args.insert(call_args.end(), {packed_args, ConstInt32(nargs), result});
  }
  auto call_callee = llvm::FunctionCallee(callee_ftype, callee_value);
  llvm::Value* call = builder_->CreateCall(call_callee, call_args);

  llvm::BasicBlock* end_block = CheckCallSuccess(call);

  PackedCall pc = {nullptr};

  auto prim_r_type = r_type.as<PrimType>();
  bool is_void = IsVoidType(r_type) || (prim_r_type && prim_r_type.value().IsVoid());
  if (!is_void) {
    // Load the return value using its packed-ABI representation, then cast it
    // to the exact type requested by the TIR call.
    llvm::Type* llvm_r_api_type = [&]() -> llvm::Type* {
      if (prim_r_type) {
        return DTypeToLLVMType(tirx::APIType(prim_r_type.value()));
      }
      TVM_FFI_ICHECK(r_type.as<PointerTypeNode>())
          << "Packed calls may return only primitive or pointer types, but got " << r_type;
      return t_void_p_;
    }();
    llvm::Value* result_value =
        builder_->CreateInBoundsGEP(t_tvm_ffi_any_, result, {ConstInt32(0), ConstInt32(2)});
    llvm::Value* load_ptr =
        builder_->CreatePointerCast(result_value, llvmGetPointerTo(llvm_r_api_type, 0));
    llvm::Value* rvalue = builder_->CreateAlignedLoad(llvm_r_api_type, load_ptr, llvm::Align(8));

    if (prim_r_type) {
      PrimType r_api_type = tirx::APIType(prim_r_type.value());
      pc.ret_value = CreateCast(r_api_type, prim_r_type.value(), rvalue);
    } else {
      pc.ret_value = builder_->CreatePointerCast(rvalue, GetLLVMType(r_type));
    }
    llvm::Value* result_type_index =
        builder_->CreateInBoundsGEP(t_tvm_ffi_any_, result, {ConstInt32(0), ConstInt32(0)});

    // Load the return type code.
    pc.ret_type_index = builder_->CreateAlignedLoad(t_int32_, result_type_index, llvm::Align(4));
  }

  pc.end_block = end_block;
  return pc;
}

llvm::Value* CodeGenCPU::CreateCallPacked(const CallNode* op) {
  TVM_FFI_ICHECK_EQ(op->args.size(), 4U);
  bool use_string_lookup = op->op.same_as(builtin::tvm_call_packed_lowered());
  PackedCall pc = MakeCallPackedLowered(op->args, op->ty, op->args[2].as<IntImmNode>()->value,
                                        op->args[3].as<IntImmNode>()->value, use_string_lookup);
  return pc.ret_value;
}

llvm::Value* CodeGenCPU::CreateCallTracePacked(const CallNode* op) {
  TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
  PackedCall pc = MakeCallPackedLowered(op->args, op->ty, op->args[2].as<IntImmNode>()->value,
                                        op->args[3].as<IntImmNode>()->value, true);
  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  // Get traced value.
  llvm::Value* traced_value = MakeValue(op->args[4]);
  // The update_block handles case when we need to update the return value.
  llvm::BasicBlock* update_block = llvm::BasicBlock::Create(*ctx, "update_block", function_);
  // The continue_block handles case when we need to return original
  // traced value.
  llvm::BasicBlock* continue_block = llvm::BasicBlock::Create(*ctx, "continue_block", function_);

  // Check the ret_type_code and create cmp instruction.
  llvm::Value* cmp = builder_->CreateICmpNE(
      pc.ret_type_index, llvm::ConstantInt::get(t_int_, ffi::TypeIndex::kTVMFFINone));
  builder_->CreateCondBr(cmp, update_block, continue_block);
  builder_->SetInsertPoint(update_block);
  builder_->CreateBr(continue_block);
  builder_->SetInsertPoint(continue_block);
  // The return value depends on from what bb we come from.
  llvm::PHINode* phi_rvalue = builder_->CreatePHI(traced_value->getType(), 2);
  phi_rvalue->addIncoming(pc.ret_value, update_block);
  phi_rvalue->addIncoming(traced_value, pc.end_block);
  return phi_rvalue;
}

llvm::Value* CodeGenCPU::RuntimeTVMFFIFunctionCall() {
  if (f_tvm_ffi_func_call_ != nullptr) return f_tvm_ffi_func_call_;
  return GetContextPtr(gv_tvm_ffi_func_call_);
}

llvm::Value* CodeGenCPU::RuntimeTVMFFIEnvModLookupFromImports() {
  if (f_tvm_ffi_env_mod_lookup_from_imports_ != nullptr) {
    return f_tvm_ffi_env_mod_lookup_from_imports_;
  }
  return GetContextPtr(gv_tvm_ffi_env_mod_lookup_from_imports_);
}

llvm::Value* CodeGenCPU::RuntimeTVMFFIHandleInitOnce() {
  if (f_tvm_ffi_handle_init_once_ != nullptr) return f_tvm_ffi_handle_init_once_;
  return GetContextPtr(gv_tvm_ffi_handle_init_once_);
}

llvm::Value* CodeGenCPU::RuntimeTVMFFIErrorSetRaisedFromCStr() {
  if (f_tvm_ffi_set_raised_by_c_str_ != nullptr) return f_tvm_ffi_set_raised_by_c_str_;
  return GetContextPtr(gv_tvm_ffi_set_last_error_c_str_);
}
llvm::Value* CodeGenCPU::RuntimeTVMFFIErrorSetRaisedFromCStrParts() {
  if (f_tvm_ffi_set_raised_from_c_str_parts_ != nullptr) {
    return f_tvm_ffi_set_raised_from_c_str_parts_;
  }
  // fallback: create external linkage
  f_tvm_ffi_set_raised_from_c_str_parts_ = llvm::Function::Create(
      ftype_tvm_ffi_error_set_raised_from_c_str_parts_, llvm::Function::ExternalLinkage,
      "TVMFFIErrorSetRaisedFromCStrParts", module_.get());
  return f_tvm_ffi_set_raised_from_c_str_parts_;
}
/*!
 * \brief Get or create a noinline LLVM helper function for raising errors from TIR assertions.
 *
 * Creates an internal LLVM function `__tvm_set_raised_<max_n>` that:
 *   1. Stack-allocates a `const char**` array of size `max_n`
 *   2. Stores each message part argument into the array
 *   3. Calls `TVMFFIErrorSetRaisedFromCStrParts(kind, parts, n)`
 *
 * The function is marked `noinline` to keep error-raising code out of the hot path.
 * Assert codegen emits a conditional branch to a block that calls this helper,
 * so the error-handling logic does not bloat the main function body.
 *
 * Results are cached in `set_raised_helpers_` keyed by `max_n`, so each distinct
 * bucket size (6, 12, 18, ...) is created at most once per module.
 *
 * \param max_n The maximum number of message part slots (rounded up to a bucket size).
 * \return The LLVM function for raising errors with up to `max_n` message parts.
 */
llvm::Function* CodeGenCPU::GetOrCreateSetRaisedHelper(int max_n) {
  auto it = set_raised_helpers_.find(max_n);
  if (it != set_raised_helpers_.end()) return it->second;

  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  // Build function type: void(i8* kind, i32 n, i8* msg0, i8* msg1, ..., i8* msg_{max_n-1})
  std::vector<llvm::Type*> param_types;
  param_types.push_back(llvmGetPointerTo(t_char_, 0));  // kind
  param_types.push_back(t_int_);                        // n
  for (int i = 0; i < max_n; ++i) {
    param_types.push_back(llvmGetPointerTo(t_char_, 0));  // msg_i
  }
  auto* ftype = llvm::FunctionType::get(t_void_, param_types, false);
  std::string fname = "__tvm_set_raised_" + std::to_string(max_n);
  auto* func = llvm::Function::Create(ftype, llvm::Function::InternalLinkage, fname, module_.get());
  func->addFnAttr(llvm::Attribute::NoInline);

  // Create the function body
  auto* entry = llvm::BasicBlock::Create(*ctx, "entry", func);
  llvm::IRBuilder<> b(entry);

  // Alloca array of char* on the stack
  auto* arr_type = llvm::ArrayType::get(llvmGetPointerTo(t_char_, 0), max_n);
  auto* parts_alloca = b.CreateAlloca(arr_type);

  // Store each msg_i into parts[i]
  auto arg_it = func->arg_begin();
  llvm::Value* kind_arg = &*arg_it++;
  llvm::Value* n_arg = &*arg_it++;
  for (int i = 0; i < max_n; ++i) {
    llvm::Value* msg_i = &*arg_it++;
    llvm::Value* gep = b.CreateInBoundsGEP(arr_type, parts_alloca, {ConstInt32(0), ConstInt32(i)});
    b.CreateStore(msg_i, gep);
  }

  // Cast alloca to const char**
  llvm::Value* parts_ptr =
      b.CreateInBoundsGEP(arr_type, parts_alloca, {ConstInt32(0), ConstInt32(0)});

  // Call TVMFFIErrorSetRaisedFromCStrParts(kind, parts, n)
  auto callee = llvm::FunctionCallee(ftype_tvm_ffi_error_set_raised_from_c_str_parts_,
                                     RuntimeTVMFFIErrorSetRaisedFromCStrParts());
  b.CreateCall(callee, {kind_arg, parts_ptr, n_arg});
  b.CreateRetVoid();

  set_raised_helpers_[max_n] = func;
  return func;
}
llvm::Value* CodeGenCPU::RuntimeTVMParallelLaunch() {
  if (f_tvm_parallel_launch_ != nullptr) return f_tvm_parallel_launch_;
  return GetContextPtr(gv_tvm_parallel_launch_);
}

llvm::Value* CodeGenCPU::RuntimeTVMParallelBarrier() {
  if (f_tvm_parallel_barrier_ != nullptr) return f_tvm_parallel_barrier_;
  return GetContextPtr(gv_tvm_parallel_barrier_);
}

void CodeGenCPU::AddStartupFunction() {
  if (!target_c_runtime_) {
    llvm::FunctionType* ftype = llvm::FunctionType::get(t_void_, {}, false);
    function_ = llvm::Function::Create(ftype, llvm::Function::InternalLinkage,
                                       "__tvm_module_startup", module_.get());
    SetTargetAttributes(function_);
    llvm::BasicBlock* startup_entry =
        llvm::BasicBlock::Create(*llvm_target_->GetContext(), "entry", function_);
    builder_->SetInsertPoint(startup_entry);
    for (const auto& kv : export_system_symbols_) {
      llvm::Value* name = GetConstString(kv.first);
      builder_->CreateCall(f_tvm_register_system_symbol_,
                           {name, builder_->CreateBitCast(kv.second, t_void_p_)});
    }
    llvm::appendToGlobalCtors(*module_, function_, 65535);
    builder_->CreateRet(nullptr);
  }
}

llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
  const ffi::Array<Expr>& args = op->args;
  if (op->op.same_as(builtin::tvm_call_packed_lowered())) {
    return CreateCallPacked(op);
  } else if (op->op.same_as(builtin::tvm_call_trace_packed_lowered())) {
    return CreateCallTracePacked(op);
  } else if (op->op.same_as(builtin::tvm_call_cpacked_lowered())) {
    return CreateCallPacked(op);
  } else if (op->op.same_as(builtin::tvm_static_handle())) {
    return CreateStaticHandle();
  } else if (op->op.same_as(builtin::tvm_throw_last_error())) {
    builder_->CreateRet(ConstInt32(-1));
    auto next_block = std::next(builder_->GetInsertBlock()->getIterator());
    llvm::BasicBlock* new_bb =
        llvm::BasicBlock::Create(*llvm_target_->GetContext(), "cont", function_, &*next_block);
    builder_->SetInsertPoint(new_bb);
    return ConstInt32(-1);
  } else if (op->op.same_as(builtin::tvm_struct_get())) {
    TVM_FFI_ICHECK_EQ(args.size(), 3U);
    int kind = args[2].as<IntImm>().value()->value;
    Type op_type = op->ty;
    TypedPointer ref = CreateStructRefPtr(op_type, MakeValue(args[0]), MakeValue(args[1]), kind);
    if (kind == builtin::kDLTensorAddr) {
      TVM_FFI_ICHECK(op_type.as<PointerTypeNode>())
          << "The address of a DLTensor must have pointer type, but got " << op_type;
      return builder_->CreatePointerCast(ref.addr, GetLLVMType(op_type));
    }

    llvm::Value* struct_value = builder_->CreateLoad(ref.type, ref.addr);

    if (auto prim_type = op_type.as<PrimType>()) {
      if (prim_type.value() == PrimType::Bool()) {
        struct_value = CreateCast(PrimType::Int(64), prim_type.value(), struct_value);
      }
    } else {
      TVM_FFI_ICHECK(op_type.as<PointerTypeNode>())
          << "Struct fields must have primitive or pointer type, but got " << op_type;
      llvm::Type* target = GetLLVMType(op_type);
      TVM_FFI_ICHECK(struct_value->getType()->isPointerTy());
      if (struct_value->getType() != target) {
        struct_value = builder_->CreatePointerCast(struct_value, target);
      }
    }

    return struct_value;
  } else if (op->op.same_as(builtin::tvm_struct_set())) {
    TVM_FFI_ICHECK_EQ(args.size(), 4U);
    int kind = args[2].as<IntImm>().value()->value;
    llvm::Value* value = MakeValue(args[3]);
    TypedPointer ref =
        CreateStructRefPtr(args[3]->ty, MakeValue(args[0]), MakeValue(args[1]), kind);
    TVM_FFI_ICHECK(kind != builtin::kDLTensorAddr);
    if (value->getType()->isPointerTy()) {
      value = builder_->CreatePointerCast(value, ref.type);
    }

    if (kind == builtin::kTVMFFIAnyUnionValue) {
      // when we set any union value, we need to be careful to
      // clear off the union value to zero if the set size is less than 64 bits
      if (data_layout_->getTypeAllocSize(ref.type) != 8) {
        llvm::Value* i64_addr =
            builder_->CreatePointerCast(ref.addr, llvmGetPointerTo(t_int64_, 0));
        builder_->CreateStore(ConstInt64(0), i64_addr);
      }
    }
    builder_->CreateStore(value, ref.addr);
    return ConstInt32(0);
  } else if (op->op.same_as(builtin::tvm_stack_alloca())) {
    TVM_FFI_ICHECK_EQ(args.size(), 2U);
    std::string type = args[0].as<StringImm>().value()->value;
    return WithFunctionEntry([&]() -> llvm::AllocaInst* {
      const int64_t* pval = as_const_int(args[1].as_or_throw<PrimExpr>());
      TVM_FFI_ICHECK(pval) << "require stack alloca to contain constant value";
      llvm::Value* num = ConstInt32(pval[0]);
      if (type == "shape") {
        return builder_->CreateAlloca(t_tvm_shape_index_, num);
      } else if (type == "tvm_ffi_any") {
        return builder_->CreateAlloca(t_tvm_ffi_any_, num);
      } else if (type == "array") {
        return builder_->CreateAlloca(t_tvm_array_, num);
      } else if (type == "tensormap") {
        auto* alloca = builder_->CreateAlloca(t_tvm_tensormap_, num);
        alloca->setAlignment(llvm::Align(64));
        return alloca;
      } else {
        TVM_FFI_THROW(InternalError) << "Unknown stack alloca type " << type;
      }
    });
  } else {
    return CodeGenLLVM::CreateIntrinsic(op);
  }
}

void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
  EmitDebugLocation(op);
  llvm::Value* cond = MakeValue(op->condition);
  llvm::LLVMContext* ctx = llvm_target_->GetContext();
  auto* fail_block = llvm::BasicBlock::Create(*ctx, "assert_fail", function_);
  auto* end_block = llvm::BasicBlock::Create(*ctx, "assert_end", function_);
  builder_->CreateCondBr(cond, end_block, fail_block, md_very_likely_branch_);
  // fail condition.
  builder_->SetInsertPoint(fail_block);

  int num_parts = static_cast<int>(op->message_parts.size());
  // Round up to the nearest helper bucket: 6, 12, 18, ...
  int max_n = ((num_parts + 5) / 6) * 6;
  if (max_n == 0) max_n = 6;

  llvm::Function* helper = GetOrCreateSetRaisedHelper(max_n);
  // Build arguments: kind, n, msg0, msg1, ..., msg_{max_n-1}
  std::vector<llvm::Value*> call_args;
  call_args.push_back(GetConstString(op->error_kind->value));
  call_args.push_back(ConstInt32(num_parts));
  for (int i = 0; i < max_n; ++i) {
    if (i < num_parts) {
      call_args.push_back(GetConstString(op->message_parts[i]->value));
    } else {
      // Pad with null
      call_args.push_back(llvm::ConstantPointerNull::get(
          llvm::cast<llvm::PointerType>(llvmGetPointerTo(t_char_, 0))));
    }
  }
  builder_->CreateCall(helper, call_args);
  builder_->CreateRet(ConstInt32(-1));
  // otherwise set it to be new end.
  builder_->SetInsertPoint(end_block);
  CodeGenLLVM::VisitStmt_(op);
}

void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
  EmitDebugLocation(op);
  if (op->attr_key == tirx::attr::compute_scope) {
    this->CreateComputeScope(op);
  } else if (tirx::attr::IsPragmaKey(op->attr_key)) {
    if (op->attr_key == "pragma_parallel_stride_pattern") {
      TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
          << "Pragma parallel_stride_pattern only valid in parallel launch";
      parallel_env_.stride_pattern = true;
      this->VisitStmt(op->body);
    } else if (op->attr_key == "pragma_parallel_launch_point") {
      CreateParallelLaunch(op->body, 0, "pragma_parallel");
    } else if (op->attr_key == "pragma_parallel_barrier_when_finish") {
      TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
          << "Cannot run barrier without parallel environment";
      TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
          << "Cannot not place within parallel loop as the workload may differ, "
          << " place it between parallel and parallel_launch_point";
      this->VisitStmt(op->body);
      auto bar_callee =
          llvm::FunctionCallee(ftype_tvm_parallel_barrier_, RuntimeTVMParallelBarrier());
      builder_->CreateCall(bar_callee, {MakeValue(parallel_env_.task_id), parallel_env_.penv});
    } else if (op->attr_key == tirx::attr::pragma_import_llvm) {
      const StringImmNode* value = op->value.as<StringImmNode>();
      TVM_FFI_ICHECK(value != nullptr);
      this->HandleImport(value->value);
      this->VisitStmt(op->body);
    } else {
      LOG(WARNING) << "Unknown pragma " << op->attr_key;
      this->VisitStmt(op->body);
    }
  } else {
    CodeGenLLVM::VisitStmt_(op);
  }
}

void CodeGenCPU::VisitStmt_(const ForNode* op) {
  EmitDebugLocation(op);
  if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
    CodeGenLLVM::VisitStmt_(op);
  } else if (op->kind == ForKind::kParallel) {
    TVM_FFI_ICHECK(is_zero(op->min))
        << "Parallel launch require canonical loop with zero start index";
    TVM_FFI_ICHECK(op->HasTrivialStep())
        << "Parallel launch require canonical loop with trivial loop step";
    if (parallel_env_.penv == nullptr) {
      auto copy_node = For(ffi::make_object<ForNode>(*op));
      CreateParallelLaunch(copy_node, 0,
                           std::string("loop_parallel_") + op->loop_var->name_hint.c_str());
    } else {
      // already in parallel env.
      TVM_FFI_ICHECK(parallel_env_.task_id.defined());
      TVM_FFI_ICHECK(parallel_env_.num_task.defined());
      TVM_FFI_ICHECK(parallel_env_.penv != nullptr);
      PrimType t(op->extent.ty()->dtype);
      PrimExpr num_task = cast(t, parallel_env_.num_task.as_or_throw<PrimExpr>());
      PrimExpr task_id = cast(t, parallel_env_.task_id.as_or_throw<PrimExpr>());
      TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
          << "Nested parallel loop is not supported by threadpool, try fuse them instead";
      parallel_env_.in_parallel_loop = true;
      PrimExpr end = is_zero(op->min) ? op->extent : analyzer_->Simplify(op->min + op->extent);
      if (parallel_env_.stride_pattern) {
        CreateSerialFor(MakeValue(task_id), MakeValue(end), MakeValue(num_task), op->loop_var,
                        op->body);
      } else {
        PrimExpr step = (op->extent + num_task - IntImm(t, 1)) / num_task;
        PrimExpr begin = min(task_id * step, op->extent);
        end = min((task_id + IntImm(t, 1)) * step, end);
        CreateSerialFor(MakeValue(begin), MakeValue(end),
                        llvm::ConstantInt::getSigned(GetLLVMType(end), 1), op->loop_var, op->body);
      }
      parallel_env_.in_parallel_loop = false;
      ++parallel_env_.parallel_loop_count;
    }
  } else {
    TVM_FFI_THROW(InternalError) << "cannot handle for type " << op->kind;
  }
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def_packed("tvm.codegen.llvm.target_cpu",
                               [](const ffi::PackedArgs& targs, ffi::Any* rv) {
                                 *rv = static_cast<void*>(new CodeGenCPU());
                               });
}

}  // namespace codegen
}  // namespace tvm

#endif  // TVM_LLVM_VERSION
