[Mlir-commits] [flang] [llvm] [mlir] [openmp] [Flang][OpenMP] Lower scan directive and inscan reduction modifier (PR #206747)

[CHANDRA GHALE]

================
@@ -5661,17 +5703,67 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitScanReduction(
         ScanRedInfo->Span,
         llvm::ConstantInt::get(ScanRedInfo->Span->getType(), 1));
     Builder.SetInsertPoint(InputBB);
-    Builder.CreateBr(LoopBB);
+    // Branch around the orig-val seed load and combine when the loop has zero
+    // logical iterations. When Span == 0 no scan-buffer entry is populated, so
+    // loading buffer[0] and feeding it to a (possibly input-sensitive
+    // user-defined) reduction combiner would be undefined behavior. In that
+    // case skip straight to the exit; the finals' zero-trip guard preserves
+    // orig-val as the result.
+    llvm::Value *HasElem = Builder.CreateICmpUGT(
+        ScanRedInfo->Span, llvm::ConstantInt::get(IndexTy, 0));
+    llvm::BasicBlock *SeedBB =
+        BasicBlock::Create(CurFn->getContext(), "omp.scan.seed", CurFn, ExitBB);
+    Builder.CreateCondBr(HasElem, SeedBB, ExitBB);
+    Builder.SetInsertPoint(SeedBB);
+    // Combine the original variable's incoming value (orig-val) into the first
+    // buffer element before computing the prefix sum. Per the OpenMP scan
+    // semantics orig-val is a single prefix element that must be reflected in
+    // every inclusive/exclusive scan result and in the final reduction value.
+    // Folding it in here lets the log-scan prefix computation propagate it into
+    // every element (and into the finals' read of `buffer[Span]`), regardless
+    // of whether an iteration's input phase accumulates into or overwrites the
+    // reduction variable. Span > 0 is guaranteed here, so buffer[1] is valid.
+    for (ReductionInfo RedInfo : ReductionInfos) {
+      Value *ReductionVal = RedInfo.PrivateVariable;
+      Value *BuffPtr = (*(ScanRedInfo->ScanBuffPtrs))[ReductionVal];
+      Value *Buff = Builder.CreateLoad(Builder.getPtrTy(), BuffPtr);
+      Type *DestTy = RedInfo.ElementType;
+      Value *ElemPtr = Builder.CreateInBoundsGEP(
+          DestTy, Buff, llvm::ConstantInt::get(IndexTy, 1), "arrayOffset");
+      Value *OrigVal = Builder.CreateLoad(DestTy, RedInfo.Variable);
+      Value *Elem = Builder.CreateLoad(DestTy, ElemPtr);
+      llvm::Value *Combined;
+      // orig-val is the leftmost/earliest element of the prefix, so it must be
+      // the accumulator (`omp_out`, the combiner's first/lhs operand) while the
+      // first buffer element is the incoming value (`omp_in`). The combiner is
+      // associative but not necessarily commutative, so this order matters.
+      InsertPointOrErrorTy AfterIP =
+          RedInfo.ReductionGen(Builder.saveIP(), OrigVal, Elem, Combined);
+      if (!AfterIP)
+        return AfterIP.takeError();
+      Builder.restoreIP(*AfterIP);
+      Builder.CreateStore(Combined, ElemPtr);
+    }
+    // The log-scan prefix computation is only well-defined when there are at
+    // least two elements: `log2(Span)` is invalid for Span == 0 and the loop
+    // trip computation underflows/loops forever for Span <= 1. For Span == 1
+    // the buffer already holds the final (seeded) value, so skip straight to
+    // the exit.
+    llvm::Value *SpanGuard = Builder.CreateICmpUGT(
+        ScanRedInfo->Span,
+        llvm::ConstantInt::get(ScanRedInfo->Span->getType(), 1));
+    Builder.CreateCondBr(SpanGuard, LoopBB, ExitBB);
     emitBlock(LoopBB, CurFn);
     Builder.SetInsertPoint(LoopBB);
 
     PHINode *Counter = Builder.CreatePHI(Builder.getInt32Ty(), 2);
     // size pow2k = 1;
-    PHINode *Pow2K = Builder.CreatePHI(Builder.getInt32Ty(), 2);
+    PHINode *Pow2K = Builder.CreatePHI(IndexTy, 2);
+    // The loop is now entered from the seed block (guarded by Span > 0), not
+    // directly from InputBB.
     Counter->addIncoming(llvm::ConstantInt::get(Builder.getInt32Ty(), 0),
-                         InputBB);
-    Pow2K->addIncoming(llvm::ConstantInt::get(Builder.getInt32Ty(), 1),
-                       InputBB);
+                         SeedBB);
+    Pow2K->addIncoming(llvm::ConstantInt::get(IndexTy, 1), SeedBB);
----------------
chandraghale wrote:

Thanks. The combiner can span multiple blocks, so I now capture the block the builder is actually in right before the branch into the loop and use that as the PHI predecessor, instead of hardcoding SeedBB. That way a multi-block combiner points the loop PHIs at the correct incoming block.

https://github.com/llvm/llvm-project/pull/206747