//===------ LoopGeneratorsKMP.cpp - IR helper to create loops -------------===//

//

// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.

// See https://llvm.org/LICENSE.txt for license information.

// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

//

//===----------------------------------------------------------------------===//

//

// This file contains functions to create parallel loops as LLVM-IR.

//

//===----------------------------------------------------------------------===//

#include "polly/CodeGen/LoopGeneratorsKMP.h"

#include "llvm/IR/Dominators.h"

#include "llvm/IR/Module.h"

using namespace llvm;

using namespace polly;

void ParallelLoopGeneratorKMP::createCallSpawnThreads(Value *SubFn,

                                                      Value *SubFnParam,

                                                      Value *LB, Value *UB,

                                                      Value *Stride) {

  const std::string Name = "__kmpc_fork_call";

  Function *F = M->getFunction(Name);

  Type *KMPCMicroTy = M->getTypeByName("kmpc_micro");

  if (!KMPCMicroTy) {

    // void (*kmpc_micro)(kmp_int32 *global_tid, kmp_int32 *bound_tid, ...)

    Type *MicroParams[] = {Builder.getInt32Ty()->getPointerTo(),

                           Builder.getInt32Ty()->getPointerTo()};

    KMPCMicroTy = FunctionType::get(Builder.getVoidTy(), MicroParams, true);

  }

  // If F is not available, declare it.

  if (!F) {

    StructType *IdentTy = M->getTypeByName("struct.ident_t");

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(), Builder.getInt32Ty(),

                      KMPCMicroTy->getPointerTo()};

    FunctionType *Ty = FunctionType::get(Builder.getVoidTy(), Params, true);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  Value *Task = Builder.CreatePointerBitCastOrAddrSpaceCast(

      SubFn, KMPCMicroTy->getPointerTo());

  Value *Args[] = {SourceLocationInfo,

                   Builder.getInt32(4) /* Number of arguments (w/o Task) */,

                   Task,

                   LB,

                   UB,

                   Stride,

                   SubFnParam};

  Builder.CreateCall(F, Args);

}

void ParallelLoopGeneratorKMP::deployParallelExecution(Value *SubFn,

                                                       Value *SubFnParam,

                                                       Value *LB, Value *UB,

                                                       Value *Stride) {

  // Inform OpenMP runtime about the number of threads if greater than zero

  if (PollyNumThreads > 0) {

    Value *GlobalThreadID = createCallGlobalThreadNum();

    createCallPushNumThreads(GlobalThreadID, Builder.getInt32(PollyNumThreads));

  }

  // Tell the runtime we start a parallel loop

  createCallSpawnThreads(SubFn, SubFnParam, LB, UB, Stride);

}

Function *ParallelLoopGeneratorKMP::prepareSubFnDefinition(Function *F) const {

  std::vector<Type *> Arguments = {Builder.getInt32Ty()->getPointerTo(),

                                   Builder.getInt32Ty()->getPointerTo(),

                                   LongType,

                                   LongType,

                                   LongType,

                                   Builder.getInt8PtrTy()};

  FunctionType *FT = FunctionType::get(Builder.getVoidTy(), Arguments, false);

  Function *SubFn = Function::Create(FT, Function::InternalLinkage,

                                     F->getName() + "_polly_subfn", M);

  // Name the function's arguments

  Function::arg_iterator AI = SubFn->arg_begin();

  AI->setName("polly.kmpc.global_tid");

  std::advance(AI, 1);

  AI->setName("polly.kmpc.bound_tid");

  std::advance(AI, 1);

  AI->setName("polly.kmpc.lb");

  std::advance(AI, 1);

  AI->setName("polly.kmpc.ub");

  std::advance(AI, 1);

  AI->setName("polly.kmpc.inc");

  std::advance(AI, 1);

  AI->setName("polly.kmpc.shared");

  return SubFn;

}

// Create a subfunction of the following (preliminary) structure:

//

//    PrevBB

//       |

//       v

//    HeaderBB

//       |   _____

//       v  v    |

//   CheckNextBB  PreHeaderBB

//       |\       |

//       | \______/

//       |

//       v

//     ExitBB

//

// HeaderBB will hold allocations, loading of variables and kmp-init calls.

// CheckNextBB will check for more work (dynamic) or will be "empty" (static).

// If there is more work to do: go to PreHeaderBB, otherwise go to ExitBB.

// PreHeaderBB loads the new boundaries (& will lead to the loop body later on).

// Just like CheckNextBB: PreHeaderBB is empty in the static scheduling case.

// ExitBB marks the end of the parallel execution.

// The possibly empty BasicBlocks will automatically be removed.

std::tuple<Value *, Function *>

ParallelLoopGeneratorKMP::createSubFn(Value *StrideNotUsed,

                                      AllocaInst *StructData,

                                      SetVector<Value *> Data, ValueMapT &Map) {

  Function *SubFn = createSubFnDefinition();

  LLVMContext &Context = SubFn->getContext();

  // Store the previous basic block.

  BasicBlock *PrevBB = Builder.GetInsertBlock();

  // Create basic blocks.

  BasicBlock *HeaderBB = BasicBlock::Create(Context, "polly.par.setup", SubFn);

  BasicBlock *ExitBB = BasicBlock::Create(Context, "polly.par.exit", SubFn);

  BasicBlock *CheckNextBB =

      BasicBlock::Create(Context, "polly.par.checkNext", SubFn);

  BasicBlock *PreHeaderBB =

      BasicBlock::Create(Context, "polly.par.loadIVBounds", SubFn);

  DT.addNewBlock(HeaderBB, PrevBB);

  DT.addNewBlock(ExitBB, HeaderBB);

  DT.addNewBlock(CheckNextBB, HeaderBB);

  DT.addNewBlock(PreHeaderBB, HeaderBB);

  // Fill up basic block HeaderBB.

  Builder.SetInsertPoint(HeaderBB);

  Value *LBPtr = Builder.CreateAlloca(LongType, nullptr, "polly.par.LBPtr");

  Value *UBPtr = Builder.CreateAlloca(LongType, nullptr, "polly.par.UBPtr");

  Value *IsLastPtr = Builder.CreateAlloca(Builder.getInt32Ty(), nullptr,

                                          "polly.par.lastIterPtr");

  Value *StridePtr =

      Builder.CreateAlloca(LongType, nullptr, "polly.par.StridePtr");

  // Get iterator for retrieving the previously defined parameters.

  Function::arg_iterator AI = SubFn->arg_begin();

  // First argument holds "global thread ID".

  Value *IDPtr = &*AI;

  // Skip "bound thread ID" since it is not used (but had to be defined).

  std::advance(AI, 2);

  // Move iterator to: LB, UB, Stride, Shared variable struct.

  Value *LB = &*AI;

  std::advance(AI, 1);

  Value *UB = &*AI;

  std::advance(AI, 1);

  Value *Stride = &*AI;

  std::advance(AI, 1);

  Value *Shared = &*AI;

  Value *UserContext = Builder.CreateBitCast(Shared, StructData->getType(),

                                             "polly.par.userContext");

  extractValuesFromStruct(Data, StructData->getAllocatedType(), UserContext,

                          Map);

  const int Alignment = (is64BitArch()) ? 8 : 
40
;

  Value *ID =

      Builder.CreateAlignedLoad(IDPtr, Alignment, "polly.par.global_tid");

  Builder.CreateAlignedStore(LB, LBPtr, Alignment);

  Builder.CreateAlignedStore(UB, UBPtr, Alignment);

  Builder.CreateAlignedStore(Builder.getInt32(0), IsLastPtr, Alignment);

  Builder.CreateAlignedStore(Stride, StridePtr, Alignment);

  // Subtract one as the upper bound provided by openmp is a < comparison

  // whereas the codegenForSequential function creates a <= comparison.

  Value *AdjustedUB = Builder.CreateAdd(UB, ConstantInt::get(LongType, -1),

                                        "polly.indvar.UBAdjusted");

  Value *ChunkSize =

      ConstantInt::get(LongType, std::max<int>(PollyChunkSize, 1));

  switch (PollyScheduling) {

  case OMPGeneralSchedulingType::Dynamic:

  case OMPGeneralSchedulingType::Guided:

  case OMPGeneralSchedulingType::Runtime:

    // "DYNAMIC" scheduling types are handled below (including 'runtime')

    {

      UB = AdjustedUB;

      createCallDispatchInit(ID, LB, UB, Stride, ChunkSize);

      Value *HasWork =

          createCallDispatchNext(ID, IsLastPtr, LBPtr, UBPtr, StridePtr);

      Value *HasIteration =

          Builder.CreateICmp(llvm::CmpInst::Predicate::ICMP_EQ, HasWork,

                             Builder.getInt32(1), "polly.hasIteration");

      Builder.CreateCondBr(HasIteration, PreHeaderBB, ExitBB);

      Builder.SetInsertPoint(CheckNextBB);

      HasWork = createCallDispatchNext(ID, IsLastPtr, LBPtr, UBPtr, StridePtr);

      HasIteration =

          Builder.CreateICmp(llvm::CmpInst::Predicate::ICMP_EQ, HasWork,

                             Builder.getInt32(1), "polly.hasWork");

      Builder.CreateCondBr(HasIteration, PreHeaderBB, ExitBB);

      Builder.SetInsertPoint(PreHeaderBB);

      LB = Builder.CreateAlignedLoad(LBPtr, Alignment, "polly.indvar.LB");

      UB = Builder.CreateAlignedLoad(UBPtr, Alignment, "polly.indvar.UB");

    }

    break;

  case OMPGeneralSchedulingType::StaticChunked:

  case OMPGeneralSchedulingType::StaticNonChunked:

    // "STATIC" scheduling types are handled below

    {

      createCallStaticInit(ID, IsLastPtr, LBPtr, UBPtr, StridePtr, ChunkSize);

      LB = Builder.CreateAlignedLoad(LBPtr, Alignment, "polly.indvar.LB");

      UB = Builder.CreateAlignedLoad(UBPtr, Alignment, "polly.indvar.UB");

      Value *AdjUBOutOfBounds =

          Builder.CreateICmp(llvm::CmpInst::Predicate::ICMP_SLT, UB, AdjustedUB,

                             "polly.adjustedUBOutOfBounds");

      UB = Builder.CreateSelect(AdjUBOutOfBounds, UB, AdjustedUB);

      Builder.CreateAlignedStore(UB, UBPtr, Alignment);

      Value *HasIteration = Builder.CreateICmp(

          llvm::CmpInst::Predicate::ICMP_SLE, LB, UB, "polly.hasIteration");

      Builder.CreateCondBr(HasIteration, PreHeaderBB, ExitBB);

      Builder.SetInsertPoint(CheckNextBB);

      Builder.CreateBr(ExitBB);

      Builder.SetInsertPoint(PreHeaderBB);

    }

    break;

  }

  Builder.CreateBr(CheckNextBB);

  Builder.SetInsertPoint(&*--Builder.GetInsertPoint());

  BasicBlock *AfterBB;

  Value *IV = createLoop(LB, UB, Stride, Builder, LI, DT, AfterBB,

                         ICmpInst::ICMP_SLE, nullptr, true,

                         /* UseGuard */ false);

  BasicBlock::iterator LoopBody = Builder.GetInsertPoint();

  // Add code to terminate this subfunction.

  Builder.SetInsertPoint(ExitBB);

  // Static (i.e. non-dynamic) scheduling types, are terminated with a fini-call

  if (PollyScheduling == OMPGeneralSchedulingType::StaticChunked) {

    createCallStaticFini(ID);

  }

  Builder.CreateRetVoid();

  Builder.SetInsertPoint(&*LoopBody);

  return std::make_tuple(IV, SubFn);

}

Value *ParallelLoopGeneratorKMP::createCallGlobalThreadNum() {

  const std::string Name = "__kmpc_global_thread_num";

  Function *F = M->getFunction(Name);

  // If F is not available, declare it.

  if (!F) {

    StructType *IdentTy = M->getTypeByName("struct.ident_t");

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo()};

    FunctionType *Ty = FunctionType::get(Builder.getInt32Ty(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  return Builder.CreateCall(F, {SourceLocationInfo});

}

void ParallelLoopGeneratorKMP::createCallPushNumThreads(Value *GlobalThreadID,

                                                        Value *NumThreads) {

  const std::string Name = "__kmpc_push_num_threads";

  Function *F = M->getFunction(Name);

  // If F is not available, declare it.

  if (!F) {

    StructType *IdentTy = M->getTypeByName("struct.ident_t");

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(), Builder.getInt32Ty(),

                      Builder.getInt32Ty()};

    FunctionType *Ty = FunctionType::get(Builder.getVoidTy(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  Value *Args[] = {SourceLocationInfo, GlobalThreadID, NumThreads};

  Builder.CreateCall(F, Args);

}

void ParallelLoopGeneratorKMP::createCallStaticInit(Value *GlobalThreadID,

                                                    Value *IsLastPtr,

                                                    Value *LBPtr, Value *UBPtr,

                                                    Value *StridePtr,

                                                    Value *ChunkSize) {

  const std::string Name =

      is64BitArch() ? "__kmpc_for_static_init_8" : 
"__kmpc_for_static_init_4"0
;

  Function *F = M->getFunction(Name);

  StructType *IdentTy = M->getTypeByName("struct.ident_t");

  // If F is not available, declare it.

  if (!F) {

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(),

                      Builder.getInt32Ty(),

                      Builder.getInt32Ty(),

                      Builder.getInt32Ty()->getPointerTo(),

                      LongType->getPointerTo(),

                      LongType->getPointerTo(),

                      LongType->getPointerTo(),

                      LongType,

                      LongType};

    FunctionType *Ty = FunctionType::get(Builder.getVoidTy(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  // The parameter 'ChunkSize' will hold strictly positive integer values,

  // regardless of PollyChunkSize's value

  Value *Args[] = {

      SourceLocationInfo,

      GlobalThreadID,

      Builder.getInt32(int(getSchedType(PollyChunkSize, PollyScheduling))),

      IsLastPtr,

      LBPtr,

      UBPtr,

      StridePtr,

      ConstantInt::get(LongType, 1),

      ChunkSize};

  Builder.CreateCall(F, Args);

}

void ParallelLoopGeneratorKMP::createCallStaticFini(Value *GlobalThreadID) {

  const std::string Name = "__kmpc_for_static_fini";

  Function *F = M->getFunction(Name);

  StructType *IdentTy = M->getTypeByName("struct.ident_t");

  // If F is not available, declare it.

  if (!F) {

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(), Builder.getInt32Ty()};

    FunctionType *Ty = FunctionType::get(Builder.getVoidTy(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  Value *Args[] = {SourceLocationInfo, GlobalThreadID};

  Builder.CreateCall(F, Args);

}

void ParallelLoopGeneratorKMP::createCallDispatchInit(Value *GlobalThreadID,

                                                      Value *LB, Value *UB,

                                                      Value *Inc,

                                                      Value *ChunkSize) {

  const std::string Name =

      is64BitArch() ? "__kmpc_dispatch_init_8" : 
"__kmpc_dispatch_init_4"0
;

  Function *F = M->getFunction(Name);

  StructType *IdentTy = M->getTypeByName("struct.ident_t");

  // If F is not available, declare it.

  if (!F) {

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(),

                      Builder.getInt32Ty(),

                      Builder.getInt32Ty(),

                      LongType,

                      LongType,

                      LongType,

                      LongType};

    FunctionType *Ty = FunctionType::get(Builder.getVoidTy(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  // The parameter 'ChunkSize' will hold strictly positive integer values,

  // regardless of PollyChunkSize's value

  Value *Args[] = {

      SourceLocationInfo,

      GlobalThreadID,

      Builder.getInt32(int(getSchedType(PollyChunkSize, PollyScheduling))),

      LB,

      UB,

      Inc,

      ChunkSize};

  Builder.CreateCall(F, Args);

}

Value *ParallelLoopGeneratorKMP::createCallDispatchNext(Value *GlobalThreadID,

                                                        Value *IsLastPtr,

                                                        Value *LBPtr,

                                                        Value *UBPtr,

                                                        Value *StridePtr) {

  const std::string Name =

      is64BitArch() ? "__kmpc_dispatch_next_8" : 
"__kmpc_dispatch_next_4"0
;

  Function *F = M->getFunction(Name);

  StructType *IdentTy = M->getTypeByName("struct.ident_t");

  // If F is not available, declare it.

  if (!F) {

    GlobalValue::LinkageTypes Linkage = Function::ExternalLinkage;

    Type *Params[] = {IdentTy->getPointerTo(),

                      Builder.getInt32Ty(),

                      Builder.getInt32Ty()->getPointerTo(),

                      LongType->getPointerTo(),

                      LongType->getPointerTo(),

                      LongType->getPointerTo()};

    FunctionType *Ty = FunctionType::get(Builder.getInt32Ty(), Params, false);

    F = Function::Create(Ty, Linkage, Name, M);

  }

  Value *Args[] = {SourceLocationInfo, GlobalThreadID, IsLastPtr, LBPtr, UBPtr,

                   StridePtr};

  return Builder.CreateCall(F, Args);

}

// TODO: This function currently creates a source location dummy. It might be

// necessary to (actually) provide information, in the future.

GlobalVariable *ParallelLoopGeneratorKMP::createSourceLocation() {

  const std::string LocName = ".loc.dummy";

  GlobalVariable *SourceLocDummy = M->getGlobalVariable(LocName);

  if (SourceLocDummy == nullptr) {

    const std::string StructName = "struct.ident_t";

    StructType *IdentTy = M->getTypeByName(StructName);

    // If the ident_t StructType is not available, declare it.

    // in LLVM-IR: ident_t = type { i32, i32, i32, i32, i8* }

    if (!IdentTy) {

      Type *LocMembers[] = {Builder.getInt32Ty(), Builder.getInt32Ty(),

                            Builder.getInt32Ty(), Builder.getInt32Ty(),

                            Builder.getInt8PtrTy()};

      IdentTy =

          StructType::create(M->getContext(), LocMembers, StructName, false);

    }

    const auto ArrayType =

        llvm::ArrayType::get(Builder.getInt8Ty(), /* Length */ 23);

    // Global Variable Definitions

    GlobalVariable *StrVar = new GlobalVariable(

        *M, ArrayType, true, GlobalValue::PrivateLinkage, 0, ".str.ident");

    StrVar->setAlignment(1);

    SourceLocDummy = new GlobalVariable(

        *M, IdentTy, true, GlobalValue::PrivateLinkage, nullptr, LocName);

    SourceLocDummy->setAlignment(8);

    // Constant Definitions

    Constant *InitStr = ConstantDataArray::getString(

        M->getContext(), "Source location dummy.", true);

    Constant *StrPtr = static_cast<Constant *>(Builder.CreateInBoundsGEP(

        ArrayType, StrVar, {Builder.getInt32(0), Builder.getInt32(0)}));

    Constant *LocInitStruct = ConstantStruct::get(

        IdentTy, {Builder.getInt32(0), Builder.getInt32(0), Builder.getInt32(0),

                  Builder.getInt32(0), StrPtr});

    // Initialize variables

    StrVar->setInitializer(InitStr);

    SourceLocDummy->setInitializer(LocInitStruct);

  }

  return SourceLocDummy;

}

bool ParallelLoopGeneratorKMP::is64BitArch() {

  return (LongType->getIntegerBitWidth() == 64);

}

OMPGeneralSchedulingType ParallelLoopGeneratorKMP::getSchedType(

    int ChunkSize, OMPGeneralSchedulingType Scheduling) const {

  if (ChunkSize == 0 && 
Scheduling == OMPGeneralSchedulingType::StaticChunked8
)

    return OMPGeneralSchedulingType::StaticNonChunked;

  return Scheduling;

}