[MLIR][OpenMP][OMPIRBuilder] Support omp.target 'if' translation to LLVM IR

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

@@ -2871,6 +2871,7 @@ class OpenMPIRBuilder {
   /// \param Loc where the target data construct was encountered.
   /// \param IsSPMD whether this is an SPMD target launch.
   /// \param IsOffloadEntry whether it is an offload entry.
+  /// \param IfCond value of the IF clause for the TARGET construct or nullptr.
   /// \param CodeGenIP The insertion point where the call to the outlined
   /// function should be emitted.
   /// \param EntryInfo The entry information about the function.
@@ -2884,7 +2885,7 @@ class OpenMPIRBuilder {
   /// \param Dependencies A vector of DependData objects that carry
   // dependency information as passed in the depend clause.
   InsertPointTy createTarget(const LocationDescription &Loc, bool IsSPMD,
-                             bool IsOffloadEntry,
+                             bool IsOffloadEntry, Value *IfCond,
                              OpenMPIRBuilder::InsertPointTy AllocaIP,
                              OpenMPIRBuilder::InsertPointTy CodeGenIP,
                              TargetRegionEntryInfo &EntryInfo,

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -7266,7 +7266,7 @@ static void emitTargetCall(
     const OpenMPIRBuilder::TargetKernelDefaultBounds &DefaultBounds,
     const OpenMPIRBuilder::TargetKernelRuntimeBounds &RuntimeBounds,
     Function *OutlinedFn, Constant *OutlinedFnID,
-    SmallVectorImpl<Value *> &Args,
+    SmallVectorImpl<Value *> &Args, Value *IfCond,
     OpenMPIRBuilder::GenMapInfoCallbackTy GenMapInfoCB,
     SmallVector<llvm::OpenMPIRBuilder::DependData> Dependencies = {}) {
   // Generate a function call to the host fallback implementation of the target
@@ -7283,9 +7283,7 @@ static void emitTargetCall(
   bool HasDependencies = Dependencies.size() > 0;
   bool RequiresOuterTargetTask = HasNoWait || HasDependencies;
 
-  // If we don't have an ID for the target region, it means an offload entry
-  // wasn't created. In this case we just run the host fallback directly.
-  if (!OutlinedFnID) {
+  auto &&EmitTargetCallElse = [&]() {
     if (RequiresOuterTargetTask) {
       // Arguments that are intended to be directly forwarded to an
       // emitKernelLaunch call are pased as nullptr, since OutlinedFnID=nullptr
@@ -7298,96 +7296,132 @@ static void emitTargetCall(
     } else {
       Builder.restoreIP(EmitTargetCallFallbackCB(Builder.saveIP()));
     }
-    return;
-  }
-
-  OpenMPIRBuilder::TargetDataInfo Info(
-      /*RequiresDevicePointerInfo=*/false,
-      /*SeparateBeginEndCalls=*/true);
-
-  OpenMPIRBuilder::MapInfosTy &MapInfo = GenMapInfoCB(Builder.saveIP());
-  OpenMPIRBuilder::TargetDataRTArgs RTArgs;
-  OMPBuilder.emitOffloadingArraysAndArgs(AllocaIP, Builder.saveIP(), Info,
-                                         RTArgs, MapInfo,
-                                         /*IsNonContiguous=*/true,
-                                         /*ForEndCall=*/false);
-
-  SmallVector<Value *, 3> NumTeamsC;
-  for (auto [DefNumTeams, RtNumTeams] :
-       llvm::zip_equal(DefaultBounds.MaxTeams, RuntimeBounds.MaxTeams)) {
-    NumTeamsC.push_back(RtNumTeams ? RtNumTeams
-                                   : Builder.getInt32(DefNumTeams));
-  }
-
-  // Calculate number of threads: 0 if no clauses specified, otherwise it is the
-  // minimum between optional THREAD_LIMIT and MAX_THREADS clauses. Perform a
-  // type cast to uint32.
-  auto InitMaxThreadsClause = [&Builder](Value *Clause) {
-    if (Clause)
-      Clause = Builder.CreateIntCast(Clause, Builder.getInt32Ty(),
-                                     /*isSigned=*/false);
-    return Clause;
   };
 
-  auto CombineMaxThreadsClauses = [&Builder](Value *Clause, Value *&Result) {
-    if (Clause)
-      Result = Result
-                   ? Builder.CreateSelect(Builder.CreateICmpULT(Result, Clause),
+  auto &&EmitTargetCallThen = [&]() {
+    OpenMPIRBuilder::TargetDataInfo Info(
+        /*RequiresDevicePointerInfo=*/false,
+        /*SeparateBeginEndCalls=*/true);
+
+    OpenMPIRBuilder::MapInfosTy &MapInfo = GenMapInfoCB(Builder.saveIP());
+    OpenMPIRBuilder::TargetDataRTArgs RTArgs;
+    OMPBuilder.emitOffloadingArraysAndArgs(AllocaIP, Builder.saveIP(), Info,
+                                           RTArgs, MapInfo,
+                                           /*IsNonContiguous=*/true,
+                                           /*ForEndCall=*/false);
+
+    SmallVector<Value *, 3> NumTeamsC;
+    for (auto [DefNumTeams, RtNumTeams] :
+         llvm::zip_equal(DefaultBounds.MaxTeams, RuntimeBounds.MaxTeams)) {
+      NumTeamsC.push_back(RtNumTeams ? RtNumTeams
+                                     : Builder.getInt32(DefNumTeams));
+    }
+
+    // Calculate number of threads: 0 if no clauses specified, otherwise it is
+    // the minimum between optional THREAD_LIMIT and MAX_THREADS clauses.
+    // Perform a type cast to uint32.
+    auto InitMaxThreadsClause = [&Builder](Value *Clause) {
+      if (Clause)
+        Clause = Builder.CreateIntCast(Clause, Builder.getInt32Ty(),
+                                       /*isSigned=*/false);
+      return Clause;
+    };
+
+    auto CombineMaxThreadsClauses = [&Builder](Value *Clause, Value *&Result) {
+      if (Clause)
+        Result =
+            Result ? Builder.CreateSelect(Builder.CreateICmpULT(Result, Clause),
                                           Result, Clause)
                    : Clause;
-  };
+    };
+
+    // TODO: Check if this is the correct handling for multi-dim thread_limit.
+    SmallVector<Value *, 3> NumThreadsC;
+    Value *MaxThreadsClause = InitMaxThreadsClause(RuntimeBounds.MaxThreads);
 
-  // TODO: Check if this is the correct handling for multi-dim thread_limit.
-  SmallVector<Value *, 3> NumThreadsC;
-  Value *MaxThreadsClause = InitMaxThreadsClause(RuntimeBounds.MaxThreads);
+    for (auto [RtTeamsThreadLimit, RtTargetThreadLimit] : llvm::zip_equal(
+             RuntimeBounds.TeamsThreadLimit, RuntimeBounds.TargetThreadLimit)) {
+      Value *TeamsThreadLimitClause = InitMaxThreadsClause(RtTeamsThreadLimit);
+      Value *NumThreads = InitMaxThreadsClause(RtTargetThreadLimit);
 
-  for (auto [RtTeamsThreadLimit, RtTargetThreadLimit] : llvm::zip_equal(
-           RuntimeBounds.TeamsThreadLimit, RuntimeBounds.TargetThreadLimit)) {
-    Value *TeamsThreadLimitClause = InitMaxThreadsClause(RtTeamsThreadLimit);
-    Value *NumThreads = InitMaxThreadsClause(RtTargetThreadLimit);
+      CombineMaxThreadsClauses(TeamsThreadLimitClause, NumThreads);
+      CombineMaxThreadsClauses(MaxThreadsClause, NumThreads);
 
-    CombineMaxThreadsClauses(TeamsThreadLimitClause, NumThreads);
-    CombineMaxThreadsClauses(MaxThreadsClause, NumThreads);
+      NumThreadsC.push_back(NumThreads ? NumThreads : Builder.getInt32(0));
+    }
 
-    NumThreadsC.push_back(NumThreads ? NumThreads : Builder.getInt32(0));
+    unsigned NumTargetItems = Info.NumberOfPtrs;
+    // TODO: Use correct device ID
+    Value *DeviceID = Builder.getInt64(OMP_DEVICEID_UNDEF);
+    uint32_t SrcLocStrSize;
+    Constant *SrcLocStr = OMPBuilder.getOrCreateDefaultSrcLocStr(SrcLocStrSize);
+    Value *RTLoc = OMPBuilder.getOrCreateIdent(SrcLocStr, SrcLocStrSize,
+                                               llvm::omp::IdentFlag(0), 0);
+
+    Value *TripCount = RuntimeBounds.LoopTripCount
+                           ? Builder.CreateIntCast(RuntimeBounds.LoopTripCount,
+                                                   Builder.getInt64Ty(),
+                                                   /*isSigned=*/false)
+                           : Builder.getInt64(0);
+
+    // TODO: Use correct DynCGGroupMem
+    Value *DynCGGroupMem = Builder.getInt32(0);
+    OpenMPIRBuilder::TargetKernelArgs KArgs(NumTargetItems, RTArgs, TripCount,
+                                            NumTeamsC, NumThreadsC,
+                                            DynCGGroupMem, HasNoWait);
+
+    // The presence of certain clauses on the target directive require the
+    // explicit generation of the target task.
+    if (RequiresOuterTargetTask) {
+      Builder.restoreIP(OMPBuilder.emitTargetTask(
+          OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID,
+          RTLoc, AllocaIP, Dependencies, HasNoWait));
+    } else {
+      Builder.restoreIP(OMPBuilder.emitKernelLaunch(
+          Builder, OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs,
+          DeviceID, RTLoc, AllocaIP));
+    }
+  };
+
+  // If we don't have an ID for the target region, it means an offload entry
+  // wasn't created. In this case we just run the host fallback directly.
+  if (!OutlinedFnID) {
+    EmitTargetCallElse();
+    return;
   }
 
-  unsigned NumTargetItems = Info.NumberOfPtrs;
-  // TODO: Use correct device ID
-  Value *DeviceID = Builder.getInt64(OMP_DEVICEID_UNDEF);
-  uint32_t SrcLocStrSize;
-  Constant *SrcLocStr = OMPBuilder.getOrCreateDefaultSrcLocStr(SrcLocStrSize);
-  Value *RTLoc = OMPBuilder.getOrCreateIdent(SrcLocStr, SrcLocStrSize,
-                                             llvm::omp::IdentFlag(0), 0);
-
-  Value *TripCount = RuntimeBounds.LoopTripCount
-                         ? Builder.CreateIntCast(RuntimeBounds.LoopTripCount,
-                                                 Builder.getInt64Ty(),
-                                                 /*isSigned=*/false)
-                         : Builder.getInt64(0);
-
-  // TODO: Use correct DynCGGroupMem
-  Value *DynCGGroupMem = Builder.getInt32(0);
-  OpenMPIRBuilder::TargetKernelArgs KArgs(NumTargetItems, RTArgs, TripCount,
-                                          NumTeamsC, NumThreadsC, DynCGGroupMem,
-                                          HasNoWait);
-
-  // The presence of certain clauses on the target directive require the
-  // explicit generation of the target task.
-  if (RequiresOuterTargetTask) {
-    Builder.restoreIP(OMPBuilder.emitTargetTask(
-        OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs, DeviceID,
-        RTLoc, AllocaIP, Dependencies, HasNoWait));
-  } else {
-    Builder.restoreIP(OMPBuilder.emitKernelLaunch(
-        Builder, OutlinedFn, OutlinedFnID, EmitTargetCallFallbackCB, KArgs,
-        DeviceID, RTLoc, AllocaIP));
+  // If there's no IF clause, only generate the kernel launch code path.
+  if (!IfCond) {
+    EmitTargetCallThen();
+    return;
   }
+
+  // Create if-else to handle IF clause.
+  llvm::BasicBlock *ThenBlock =
+      BasicBlock::Create(Builder.getContext(), "omp_if.then");
+  llvm::BasicBlock *ElseBlock =
+      BasicBlock::Create(Builder.getContext(), "omp_if.else");
+  llvm::BasicBlock *ContBlock =
+      BasicBlock::Create(Builder.getContext(), "omp_if.end");
+  Builder.CreateCondBr(IfCond, ThenBlock, ElseBlock);
+
+  Function *CurFn = Builder.GetInsertBlock()->getParent();
+
+  // Emit the 'then' code.
+  OMPBuilder.emitBlock(ThenBlock, CurFn);
+  EmitTargetCallThen();
+  OMPBuilder.emitBranch(ContBlock);
+  // Emit the 'else' code.
+  OMPBuilder.emitBlock(ElseBlock, CurFn);
+  EmitTargetCallElse();
+  OMPBuilder.emitBranch(ContBlock);
+  // Emit the continuation block.
+  OMPBuilder.emitBlock(ContBlock, CurFn, /*IsFinished=*/true);
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
     const LocationDescription &Loc, bool IsSPMD, bool IsOffloadEntry,
-    InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
+    Value *IfCond, InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
     TargetRegionEntryInfo &EntryInfo,
     const TargetKernelDefaultBounds &DefaultBounds,
     const TargetKernelRuntimeBounds &RuntimeBounds,
@@ -7415,7 +7449,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTarget(
   // that represents the target region. Do that now.
   if (!Config.isTargetDevice())
     emitTargetCall(*this, Builder, AllocaIP, DefaultBounds, RuntimeBounds,
-                   OutlinedFn, OutlinedFnID, Args, GenMapInfoCB, Dependencies);
+                   OutlinedFn, OutlinedFnID, Args, IfCond, GenMapInfoCB,
+                   Dependencies);
   return Builder.saveIP();
 }
 

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -6017,9 +6017,9 @@ TEST_F(OpenMPIRBuilderTest, TargetRegion) {
   RuntimeBounds.TeamsThreadLimit.push_back(nullptr);
   RuntimeBounds.MaxTeams.push_back(nullptr);
   Builder.restoreIP(OMPBuilder.createTarget(
-      OmpLoc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, Builder.saveIP(),
-      Builder.saveIP(), EntryInfo, DefaultBounds, RuntimeBounds, Inputs,
-      GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB));
+      OmpLoc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, /*IfCond=*/nullptr,
+      Builder.saveIP(), Builder.saveIP(), EntryInfo, DefaultBounds,
+      RuntimeBounds, Inputs, GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB));
   OMPBuilder.finalize();
   Builder.CreateRetVoid();
 
@@ -6134,9 +6134,9 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
   RuntimeBounds.TeamsThreadLimit.push_back(nullptr);
   RuntimeBounds.MaxTeams.push_back(nullptr);
   Builder.restoreIP(OMPBuilder.createTarget(
-      Loc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, EntryIP, EntryIP,
-      EntryInfo, DefaultBounds, RuntimeBounds, CapturedArgs, GenMapInfoCB,
-      BodyGenCB, SimpleArgAccessorCB));
+      Loc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, /*IfCond=*/nullptr,
+      EntryIP, EntryIP, EntryInfo, DefaultBounds, RuntimeBounds, CapturedArgs,
+      GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB));
 
   Builder.CreateRetVoid();
   OMPBuilder.finalize();
@@ -6290,9 +6290,9 @@ TEST_F(OpenMPIRBuilderTest, ConstantAllocaRaise) {
   RuntimeBounds.TeamsThreadLimit.push_back(nullptr);
   RuntimeBounds.MaxTeams.push_back(nullptr);
   Builder.restoreIP(OMPBuilder.createTarget(
-      Loc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, EntryIP, EntryIP,
-      EntryInfo, DefaultBounds, RuntimeBounds, CapturedArgs, GenMapInfoCB,
-      BodyGenCB, SimpleArgAccessorCB));
+      Loc, /*IsSPMD=*/false, /*IsOffloadEntry=*/true, /*IfCond=*/nullptr,
+      EntryIP, EntryIP, EntryInfo, DefaultBounds, RuntimeBounds, CapturedArgs,
+      GenMapInfoCB, BodyGenCB, SimpleArgAccessorCB));
 
   Builder.CreateRetVoid();
   OMPBuilder.finalize();

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

@@ -3476,10 +3476,6 @@ static bool getTargetEntryUniqueInfo(llvm::TargetRegionEntryInfo &targetInfo,
 
 static bool targetOpSupported(Operation &opInst) {
   auto targetOp = cast<omp::TargetOp>(opInst);
-  if (targetOp.getIfExpr()) {
-    opInst.emitError("If clause not yet supported");
-    return false;
-  }
 
   if (targetOp.getDevice()) {
     opInst.emitError("Device clause not yet supported");
@@ -3955,8 +3951,12 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   if (Value targetThreadLimit = targetOp.getThreadLimit())
     llvmTargetThreadLimit = moduleTranslation.lookupValue(targetThreadLimit);
 
+  llvm::Value *ifCond = nullptr;
+  if (Value targetIfCond = targetOp.getIfExpr())
+    ifCond = moduleTranslation.lookupValue(targetIfCond);
+
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createTarget(
-      ompLoc, targetOp.isTargetSPMDLoop(), isOffloadEntry, allocaIP,
+      ompLoc, targetOp.isTargetSPMDLoop(), isOffloadEntry, ifCond, allocaIP,
       builder.saveIP(), entryInfo, defaultBounds, runtimeBounds, kernelInput,
       genMapInfoCB, bodyCB, argAccessorCB, dds));