[Mlir-commits] [clang] [llvm] [mlir] [MLIR][OpenMP] Add codegen for teams reductions (PR #133310)

[Sergio Afonso]

================
@@ -4434,10 +4497,24 @@ getKmpcForStaticLoopForType(Type *Ty, OpenMPIRBuilder *OMPBuilder,
 static void createTargetLoopWorkshareCall(
     OpenMPIRBuilder *OMPBuilder, WorksharingLoopType LoopType,
     BasicBlock *InsertBlock, Value *Ident, Value *LoopBodyArg,
-    Type *ParallelTaskPtr, Value *TripCount, Function &LoopBodyFn) {
-  Type *TripCountTy = TripCount->getType();
+    Type *ParallelTaskPtr, Value *TripCountOrig, Function &LoopBodyFn) {
   Module &M = OMPBuilder->M;
   IRBuilder<> &Builder = OMPBuilder->Builder;
+  Value *TripCount = TripCountOrig;
+  // The trip count is 1 larger than it should be for GPU, this is because
+  // of how the deviceRTL functions work with clang. TODO: make the trip
+  // count consistent between both so we don't have to subtract one here.
+  if (OMPBuilder->Config.isGPU()) {
+    Builder.restoreIP({InsertBlock, std::prev(InsertBlock->end())});
+    LLVMContext &Ctx = M.getContext();
+    Type *IVTy = TripCountOrig->getType();
+    Type *InternalIVTy = IVTy->getIntegerBitWidth() <= 32
+                             ? Type::getInt32Ty(Ctx)
+                             : Type::getInt64Ty(Ctx);
+    Constant *One = ConstantInt::get(InternalIVTy, 1);
+    TripCount = Builder.CreateSub(TripCountOrig, One, "modified_trip_count");
+  }
+  Type *TripCountTy = TripCount->getType();
----------------
skatrak wrote:

You're already aware of this, but just to make sure we don't forget, this and changes to _mlir/test/Target/LLVMIR/omptarget-parallel-wsloop.mlir_, _mlir/test/Target/LLVMIR/omptarget-wsloop-collapsed.mlir_ and _mlir/test/Target/LLVMIR/omptarget-wsloop.mlir_ would have to be reverted if  #133435 is merged.

Also, I think _llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp_ would probably make more sense to set `IsGPU=true` at that point.

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