new changes

Author: anchuraj

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -1602,7 +1602,8 @@ class OpenMPIRBuilder {
 
   /// Creates the buffer needed for scan reduction.
   /// \param ScanVars Scan Variables.
-  void emitScanBasedDirectiveDeclsIR(ArrayRef<llvm::Value *> ScanVars);
+  void emitScanBasedDirectiveDeclsIR(ArrayRef<llvm::Value *> ScanVars,
+                                     ArrayRef<llvm::Type *> ScanVarsType);
 
   /// Copies the result back to the reduction variable.
   /// \param ReductionInfos Array type containing the ReductionOps.
@@ -2718,6 +2719,7 @@ class OpenMPIRBuilder {
   InsertPointOrErrorTy createScan(const LocationDescription &Loc,
                                   InsertPointTy AllocaIP,
                                   ArrayRef<llvm::Value *> ScanVars,
+                                  ArrayRef<llvm::Type *> ScanVarsType,
                                   bool IsInclusive);
   /// Generator for '#omp critical'
   ///

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -4002,10 +4002,11 @@ OpenMPIRBuilder::emitNoUnwindRuntimeCall(llvm::FunctionCallee Callee,
 // inclusive scans now.
 OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createScan(
     const LocationDescription &Loc, InsertPointTy AllocaIP,
-    ArrayRef<llvm::Value *> ScanVars, bool IsInclusive) {
+    ArrayRef<llvm::Value *> ScanVars, ArrayRef<llvm::Type *> ScanVarsType,
+    bool IsInclusive) {
   if (ScanInfo.OMPFirstScanLoop) {
     Builder.restoreIP(AllocaIP);
-    emitScanBasedDirectiveDeclsIR(ScanVars);
+    emitScanBasedDirectiveDeclsIR(ScanVars, ScanVarsType);
   }
   if (!updateToLocation(Loc))
     return Loc.IP;
@@ -4014,11 +4015,11 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createScan(
 
   if (ScanInfo.OMPFirstScanLoop) {
     // Emit buffer[i] = red; at the end of the input phase.
-    for (Value *ScanVar : ScanVars) {
-      Value *Buff = ScanInfo.ReductionVarToScanBuffs[ScanVar];
-      Type *DestTy = Builder.getInt32Ty(); // ScanVars[i]->getType();
+    for (int i = 0; i < ScanVars.size(); i++) {
+      Value *Buff = ScanInfo.ReductionVarToScanBuffs[ScanVars[i]];
+      Type *DestTy = ScanVarsType[i];
       Value *Val = Builder.CreateInBoundsGEP(DestTy, Buff, IV, "arrayOffset");
-      Value *Src = Builder.CreateLoad(DestTy, ScanVar);
+      Value *Src = Builder.CreateLoad(DestTy, ScanVars[i]);
       Value *Dest = Builder.CreatePointerBitCastOrAddrSpaceCast(
           Val, DestTy->getPointerTo(defaultAS));
 
@@ -4031,25 +4032,25 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createScan(
   // Initialize the private reduction variable to 0 in each iteration.
   // It is used to copy intial values to scan buffer.
   ConstantInt *Zero = ConstantInt::get(Builder.getInt32Ty(), 0);
-  for (Value *ScanVar : ScanVars) {
-    Type *DestTy = Builder.getInt32Ty(); // ScanVars[i]->getType();
+  for (int i = 0; i < ScanVars.size(); i++) {
+    Type *DestTy = ScanVarsType[i];
     Value *Dest = Builder.CreatePointerBitCastOrAddrSpaceCast(
-        ScanVar, DestTy->getPointerTo(defaultAS));
+        ScanVars[i], DestTy->getPointerTo(defaultAS));
     Builder.CreateStore(Zero, Dest);
   }
 
   if (!ScanInfo.OMPFirstScanLoop) {
     IV = ScanInfo.IV;
     // Emit red = buffer[i]; at the entrance to the scan phase.
     // TODO: if exclusive scan, the red = buffer[i-1] needs to be updated.
-    for (Value *ScanVar : ScanVars) {
-      Value *Buff = ScanInfo.ReductionVarToScanBuffs[ScanVar];
-      Type *DestTy = Builder.getInt32Ty(); // ScanVars[i]->getType();
+    for (int i = 0; i < ScanVars.size(); i++) {
+      Value *Buff = ScanInfo.ReductionVarToScanBuffs[ScanVars[i]];
+      Type *DestTy = ScanVarsType[i];
       Value *SrcPtr =
           Builder.CreateInBoundsGEP(DestTy, Buff, IV, "arrayOffset");
       Value *Src = Builder.CreateLoad(DestTy, SrcPtr);
       Value *Dest = Builder.CreatePointerBitCastOrAddrSpaceCast(
-          ScanVar, DestTy->getPointerTo(defaultAS));
+          ScanVars[i], DestTy->getPointerTo(defaultAS));
 
       Builder.CreateStore(Src, Dest);
     }
@@ -4070,29 +4071,26 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createScan(
 }
 
 void OpenMPIRBuilder::emitScanBasedDirectiveDeclsIR(
-    ArrayRef<Value *> ScanVars) {
+    ArrayRef<Value *> ScanVars, ArrayRef<Type *> ScanVarsType) {
 
   Value *AllocSpan = Builder.CreateAdd(ScanInfo.Span, Builder.getInt32(1));
-  for (Value *ScanVar : ScanVars) {
-    llvm::Value *Buff =
-        Builder.CreateAlloca(Builder.getInt32Ty(), AllocSpan, "vla");
-    ScanInfo.ReductionVarToScanBuffs[ScanVar] = Buff;
+  for (int i = 0; i < ScanVars.size(); i++) {
+    llvm::Value *Buff = Builder.CreateAlloca(ScanVarsType[i], AllocSpan, "vla");
+    ScanInfo.ReductionVarToScanBuffs[ScanVars[i]] = Buff;
   }
 }
 
 void OpenMPIRBuilder::emitScanBasedDirectiveFinalsIR(
     SmallVector<ReductionInfo> ReductionInfos) {
-  llvm::Value *OMPLast = Builder.CreateNSWAdd(
-      ScanInfo.Span,
-      llvm::ConstantInt::get(ScanInfo.Span->getType(), 1, /*isSigned=*/false));
   unsigned int DefaultAS = M.getDataLayout().getProgramAddressSpace();
   for (ReductionInfo RedInfo : ReductionInfos) {
     Value *PrivateVar = RedInfo.PrivateVariable;
     Value *OrigVar = RedInfo.Variable;
     Value *Buff = ScanInfo.ReductionVarToScanBuffs[PrivateVar];
 
     Type *SrcTy = RedInfo.ElementType;
-    Value *Val = Builder.CreateInBoundsGEP(SrcTy, Buff, OMPLast, "arrayOffset");
+    Value *Val =
+        Builder.CreateInBoundsGEP(SrcTy, Buff, ScanInfo.Span, "arrayOffset");
     Value *Src = Builder.CreateLoad(SrcTy, Val);
     Value *Dest = Builder.CreatePointerBitCastOrAddrSpaceCast(
         OrigVar, SrcTy->getPointerTo(DefaultAS));
@@ -4120,7 +4118,7 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitScanReduction(
       (llvm::Intrinsic::ID)llvm::Intrinsic::log2, Builder.getDoubleTy());
   llvm::BasicBlock *InputBB = Builder.GetInsertBlock();
   ConstantInt *One = ConstantInt::get(Builder.getInt32Ty(), 1);
-  llvm::Value *span = Builder.CreateAdd(spanDiff, One);
+  llvm::Value *span = ScanInfo.Span; // Builder.CreateAdd(spanDiff, One);
   llvm::Value *Arg = Builder.CreateUIToFP(span, Builder.getDoubleTy());
   llvm::Value *LogVal = emitNoUnwindRuntimeCall(F, Arg, "");
   F = llvm::Intrinsic::getOrInsertDeclaration(
@@ -5456,7 +5454,7 @@ OpenMPIRBuilder::tileLoops(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
   // TODO: It would be sufficient to only sink them into body of the
   // corresponding tile loop.
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> InbetweenCode;
-  for (int i = 0; i < NumLoops - 1; ++i) {
+  for (size_t i = 0; i < NumLoops - 1; ++i) {
     CanonicalLoopInfo *Surrounding = Loops[i];
     CanonicalLoopInfo *Nested = Loops[i + 1];
 
@@ -5469,7 +5467,7 @@ OpenMPIRBuilder::tileLoops(DebugLoc DL, ArrayRef<CanonicalLoopInfo *> Loops,
   Builder.SetCurrentDebugLocation(DL);
   Builder.restoreIP(OutermostLoop->getPreheaderIP());
   SmallVector<Value *, 4> FloorCount, FloorRems;
-  for (int i = 0; i < NumLoops; ++i) {
+  for (size_t i = 0; i < NumLoops; ++i) {
     Value *TileSize = TileSizes[i];
     Value *OrigTripCount = OrigTripCounts[i];
     Type *IVType = OrigTripCount->getType();

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -1440,12 +1440,14 @@ TEST_F(OpenMPIRBuilderTest, CanonicalLoopSimple) {
 
   EXPECT_EQ(&Loop->getAfter()->front(), RetInst);
 }
-void createScan(llvm::Value *scanVar, OpenMPIRBuilder &OMPBuilder,
-                IRBuilder<> &Builder, OpenMPIRBuilder::LocationDescription Loc,
+void createScan(llvm::Value *scanVar, llvm::Type *scanType,
+                OpenMPIRBuilder &OMPBuilder, IRBuilder<> &Builder,
+                OpenMPIRBuilder::LocationDescription Loc,
                 OpenMPIRBuilder::InsertPointTy &allocaIP) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
-  ASSERT_EXPECTED_INIT(InsertPointTy, retIp,
-                       OMPBuilder.createScan(Loc, allocaIP, {scanVar}, true));
+  ASSERT_EXPECTED_INIT(
+      InsertPointTy, retIp,
+      OMPBuilder.createScan(Loc, allocaIP, {scanVar}, {scanType}, true));
   Builder.restoreIP(retIp);
 }
 
@@ -5363,7 +5365,8 @@ TEST_F(OpenMPIRBuilderTest, ScanReduction) {
   auto LoopBodyGenCB = [&](InsertPointTy CodeGenIP, llvm::Value *LC) {
     NumBodiesGenerated += 1;
     Builder.restoreIP(CodeGenIP);
-    createScan(scanVar, OMPBuilder, Builder, Loc, allocaIP);
+    createScan(scanVar, Builder.getFloatTy(), OMPBuilder, Builder, Loc,
+               allocaIP);
     return Error::success();
   };
   SmallVector<CanonicalLoopInfo *> Loops;

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -47,6 +47,7 @@
 
 using namespace mlir;
 
+llvm::SmallDenseMap<llvm::Value *, llvm::Type *> ReductionVarToType;
 namespace {
 static llvm::omp::ScheduleKind
 convertToScheduleKind(std::optional<omp::ClauseScheduleKind> schedKind) {
@@ -1140,6 +1141,11 @@ initReductionVars(OP op, ArrayRef<BlockArgument> reductionArgs,
   // variables. Although this could be done after allocas, we don't want to mess
   // up with the alloca insertion point.
   for (unsigned i = 0; i < op.getNumReductionVars(); ++i) {
+
+    llvm::Type *reductionType =
+        moduleTranslation.convertType(reductionDecls[i].getType());
+    ReductionVarToType[privateReductionVariables[i]] = reductionType;
+
     SmallVector<llvm::Value *, 1> phis;
 
     // map block argument to initializer region
@@ -1213,9 +1219,11 @@ static void collectReductionInfo(
       atomicGen = owningAtomicReductionGens[i];
     llvm::Value *variable =
         moduleTranslation.lookupValue(loop.getReductionVars()[i]);
+    llvm::Type *reductionType =
+        moduleTranslation.convertType(reductionDecls[i].getType());
+    ReductionVarToType[privateReductionVariables[i]] = reductionType;
     reductionInfos.push_back(
-        {moduleTranslation.convertType(reductionDecls[i].getType()), variable,
-         privateReductionVariables[i],
+        {reductionType, variable, privateReductionVariables[i],
          /*EvaluationKind=*/llvm::OpenMPIRBuilder::EvalKind::Scalar,
          owningReductionGens[i],
          /*ReductionGenClang=*/nullptr, atomicGen});
@@ -2603,34 +2611,34 @@ convertOmpScan(Operation &opInst, llvm::IRBuilderBase &builder,
   auto scanOp = cast<omp::ScanOp>(opInst);
   bool isInclusive = scanOp.hasInclusiveVars();
   SmallVector<llvm::Value *> llvmScanVars;
+  SmallVector<llvm::Type *> llvmScanVarsType;
   mlir::OperandRange mlirScanVars = scanOp.getInclusiveVars();
   if (!isInclusive)
     mlirScanVars = scanOp.getExclusiveVars();
   for (auto val : mlirScanVars) {
     llvm::Value *llvmVal = moduleTranslation.lookupValue(val);
-
     llvmScanVars.push_back(llvmVal);
+    llvmScanVarsType.push_back(ReductionVarToType[llvmVal]);
   }
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
       findAllocaInsertPoint(builder, moduleTranslation);
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
   llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP =
       moduleTranslation.getOpenMPBuilder()->createScan(
-          ompLoc, allocaIP, llvmScanVars, isInclusive);
+          ompLoc, allocaIP, llvmScanVars, llvmScanVarsType, isInclusive);
   if (failed(handleError(afterIP, opInst)))
     return failure();
 
   builder.restoreIP(*afterIP);
 
   // TODO: The argument of LoopnestOp is stored into the index variable and this
-  // variable is used
-  //  across scan operation. However that makes the mlir
-  //  invalid.(`Intra-iteration dependences from a statement in the structured
-  //  block sequence that precede a scan directive to a statement in the
-  //  structured block sequence that follows a scan directive must not exist,
-  //  except for dependences for the list items specified in an inclusive or
-  //  exclusive clause.`). The argument of LoopNestOp need to be loaded again
-  //  after ScanOp again so mlir generated is valid.
+  // variable is used across scan operation. However that makes the mlir
+  // invalid.(`Intra-iteration dependences from a statement in the structured
+  // block sequence that precede a scan directive to a statement in the
+  // structured block sequence that follows a scan directive must not exist,
+  // except for dependences for the list items specified in an inclusive or
+  // exclusive clause.`). The argument of LoopNestOp need to be loaded again
+  // after ScanOp again so mlir generated is valid.
   auto parentOp = scanOp->getParentOp();
   auto loopOp = cast<omp::LoopNestOp>(parentOp);
   if (loopOp) {