[OpenMP] Use the OpenMPIRBuilder for `cancel` directives

[jdoerfert]

Summary

An omp cancel parallel needs to be emitted by the OpenMPIRBuilder if
the parallel was emitted by the OpenMPIRBuilder. This patch makes
this possible. The cancel logic is shared with the cancel barriers.
Testing is done via unit tests and the clang cancel_codegen.cpp file
once D70290 lands.

[merge_guards_bot]

Unit tests: fail. 61126 tests passed, 1 failed and 728 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_mutex_requirements_mutex/thread_mutex_class/try_lock.pass.cpp

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[JonChesterfield]

I like this, thanks. Very clear.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

241

The integer numbers correspond to the kmp_cancel_kind_t enum in runtime/src/kmp.h. The target offloading presently ignores this argument, but the host version has a control flow dependency on it.

I think the enum should be shared between the compiler and the runtime, or failing that, some test code should include kmp.h and check the numbers still match.

This feels like a familiar point - I've just sent an email to openmp-dev to discuss whether we can share constants between the two without copy & paste.

[jdoerfert]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

241

I move them into OMPKinds.def, we can share them from there. I'll respond to the mailing list.

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

public:
   /// \param CheckCancelFlag Flag to indicate a cancel barrier return value
   ///                        should be checked and acted upon.
   ///
   /// \returns The insertion point after the barrier.
   InsertPointTy CreateBarrier(const LocationDescription &Loc, omp::Directive DK,
                               bool ForceSimpleCall = false,
                               bool CheckCancelFlag = true);
 
+  /// Generator for '#omp cancel'
+  ///
+  /// \param Loc The location where the directive was encountered.
+  /// \param IfCondition The evaluated 'if' clause expression, if any.
+  /// \param CanceledDirective The kind of directive that is cancled.
+  ///
+  /// \returns The insertion point after the barrier.
+  InsertPointTy CreateCancel(const LocationDescription &Loc,
+                              Value *IfCondition,
+                              omp::Directive CanceledDirective);
+
   /// Generator for '#omp parallel'
   ///
   /// \param Loc The insert and source location description.
   /// \param BodyGenCB Callback that will generate the region code.
   /// \param PrivCB Callback to copy a given variable (think copy constructor).
   /// \param FiniCB Callback to finalize variable copies.
   /// \param IfCondition The evaluated 'if' clause expression, if any.
   /// \param NumThreads The evaluated 'num_threads' clause expression, if any.
private:
 
   /// Return the (LLVM-IR) string describing the source location \p Loc.
   Constant *getOrCreateSrcLocStr(const LocationDescription &Loc);
 
   /// Return an ident_t* encoding the source location \p SrcLocStr and \p Flags.
   Value *getOrCreateIdent(Constant *SrcLocStr,
                           omp::IdentFlag Flags = omp::IdentFlag(0));
 
+  /// Generate control flow and cleanup for cancellation.
+  ///
+  /// \param CancelFlag Flag indicating if the cancellation is performed.
+  /// \param CanceledDirective The kind of directive that is cancled.
+  void emitCancelationCheckImpl(Value *CancelFlag,
+                                omp::Directive CanceledDirective);
+
   /// Generate a barrier runtime call.
   ///
   /// \param Loc The location at which the request originated and is fulfilled.
   /// \param DK The directive which caused the barrier
   /// \param ForceSimpleCall Flag to force a simple (=non-cancellation) barrier.
   /// \param CheckCancelFlag Flag to indicate a cancel barrier return value
   ///                        should be checked and acted upon.
   ///

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 #ifndef OMP_RTL
 #define OMP_RTL(Enum, Str, IsVarArg, ReturnType, ...)
 #endif
 
 #define __OMP_RTL(Name, IsVarArg, ReturnType, ...)                             \
   OMP_RTL(OMPRTL_##Name, #Name, IsVarArg, ReturnType, __VA_ARGS__)
 
 __OMP_RTL(__kmpc_barrier, false, Void, IdentPtr, Int32)
+__OMP_RTL(__kmpc_cancel, false, Int32, IdentPtr, Int32, Int32)
 __OMP_RTL(__kmpc_cancel_barrier, false, Int32, IdentPtr, Int32)
 __OMP_RTL(__kmpc_global_thread_num, false, Int32, IdentPtr)
 __OMP_RTL(__kmpc_fork_call, true, Void, IdentPtr, Int32, ParallelTaskPtr)
-__OMP_RTL(__kmpc_push_num_threads, false, Void, IdentPtr, Int32, /* Int */Int32)
-__OMP_RTL(__kmpc_push_proc_bind, false, Void, IdentPtr, Int32, /* Int */Int32)
+__OMP_RTL(__kmpc_push_num_threads, false, Void, IdentPtr, Int32,
+          /* Int */ Int32)
+__OMP_RTL(__kmpc_push_proc_bind, false, Void, IdentPtr, Int32, /* Int */ Int32)
 __OMP_RTL(__kmpc_serialized_parallel, false, Void, IdentPtr, Int32)
 __OMP_RTL(__kmpc_end_serialized_parallel, false, Void, IdentPtr, Int32)
 
 __OMP_RTL(omp_get_thread_num, false, Int32, )
 
 #undef __OMP_RTL
 #undef OMP_RTL
 

 __OMP_IDENT_FLAG(BARRIER_IMPL_SINGLE, 0x0140)
 __OMP_IDENT_FLAG(BARRIER_IMPL_WORKSHARE, 0x01C0)
 
 #undef __OMP_IDENT_FLAG
 #undef OMP_IDENT_FLAG
 
 ///}
 
+/// KMP cancel kind
+///
+///{
+
+#ifndef OMP_CANCEL_KIND
+#define OMP_CANCEL_KIND(Enum, Str, DirectiveEnum, Value)
+#endif
+
+#define __OMP_CANCEL_KIND(Name, Value)                                         \
+  OMP_CANCEL_KIND(OMP_CANCEL_KIND_##Name, #Name, OMPD_##Name, Value)
+
+__OMP_CANCEL_KIND(parallel, 1)
+__OMP_CANCEL_KIND(for, 2)
+__OMP_CANCEL_KIND(sections, 3)
+__OMP_CANCEL_KIND(taskgroup, 4)
+
+#undef __OMP_CANCEL_KIND
+#undef OMP_CANCEL_KIND
+
+///}
 
 /// Proc bind kinds
 ///
 ///{
 
 #ifndef OMP_PROC_BIND_KIND
 #define OMP_PROC_BIND_KIND(Enum, Str, Value)
 #endif

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::emitBarrierImpl(const LocationDescription &Loc, Directive Kind,
   bool UseCancelBarrier =
       !ForceSimpleCall && isLastFinalizationInfoCancellable(OMPD_parallel);
 
   Value *Result = Builder.CreateCall(
       getOrCreateRuntimeFunction(UseCancelBarrier ? OMPRTL___kmpc_cancel_barrier
                                                   : OMPRTL___kmpc_barrier),
       Args);
 
-  if (UseCancelBarrier && CheckCancelFlag) {
-    // For a cancel barrier we create two new blocks.
-    BasicBlock *BB = Builder.GetInsertBlock();
-    BasicBlock *NonCancellationBlock;
-    if (Builder.GetInsertPoint() == BB->end()) {
-      // TODO: This branch will not be needed once we moved to the
-      // OpenMPIRBuilder codegen completely.
-      NonCancellationBlock = BasicBlock::Create(
-          BB->getContext(), BB->getName() + ".cont", BB->getParent());
-    } else {
-      NonCancellationBlock = SplitBlock(BB, &*Builder.GetInsertPoint());
-      BB->getTerminator()->eraseFromParent();
-      Builder.SetInsertPoint(BB);
-    }
-    BasicBlock *CancellationBlock = BasicBlock::Create(
-        BB->getContext(), BB->getName() + ".cncl", BB->getParent());
+  if (UseCancelBarrier && CheckCancelFlag)
+    emitCancelationCheckImpl(Result, OMPD_parallel);
+
+  return Builder.saveIP();
+}
+
+OpenMPIRBuilder::InsertPointTy
+OpenMPIRBuilder::CreateCancel(const LocationDescription &Loc,
+                              Value *IfCondition,
+                              omp::Directive CanceledDirective) {
+  if (!updateToLocation(Loc))
+    return Loc.IP;
+
+  // LLVM utilities like blocks with terminators.
+  auto *UI = Builder.CreateUnreachable();
+
+  Instruction *ThenTI = UI, *ElseTI = nullptr;
+  if (IfCondition)
+    SplitBlockAndInsertIfThenElse(IfCondition, UI, &ThenTI, &ElseTI);
+  Builder.SetInsertPoint(ThenTI);
+
+  Value *CancelKind = nullptr;
+  switch (CanceledDirective) {
+#define OMP_CANCEL_KIND(Enum, Str, DirectiveEnum, Value)                       \
+  case DirectiveEnum:                                                          \
+    CancelKind = Builder.getInt32(Value);                                      \
+    break;
+#include "llvm/Frontend/OpenMP/OMPKinds.def"
+  default:
+    llvm_unreachable("Unknown cancel kind!");
+  }
+
+  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc);
+  Value *Ident = getOrCreateIdent(SrcLocStr);
+  Value *Args[] = {Ident, getOrCreateThreadID(Ident), CancelKind};
+  Value *Result = Builder.CreateCall(
+      getOrCreateRuntimeFunction(OMPRTL___kmpc_cancel), Args);
+
+  // The actual cancel logic is shared with others, e.g., cancel_barriers.
+  emitCancelationCheckImpl(Result, CanceledDirective);
+
+  // Update the insertion point and remove the terminator we introduced.
+  Builder.SetInsertPoint(UI->getParent());
+  UI->eraseFromParent();
+
+  return Builder.saveIP();
+}
+
+void OpenMPIRBuilder::emitCancelationCheckImpl(
+    Value *CancelFlag, omp::Directive CanceledDirective) {
+  assert(isLastFinalizationInfoCancellable(CanceledDirective) &&
+         "Unexpected cancellation!");
+
 
-    // Jump to them based on the return value.
-    Value *Cmp = Builder.CreateIsNull(Result);
-    Builder.CreateCondBr(Cmp, NonCancellationBlock, CancellationBlock,
-                         /* TODO weight */ nullptr, nullptr);
+  Value *Cmp = Builder.CreateIsNull(CancelFlag);
 
-    // From the cancellation block we finalize all variables and go to the
-    // post finalization block that is known to the FiniCB callback.
-    Builder.SetInsertPoint(CancellationBlock);
-    auto &FI = FinalizationStack.back();
-    FI.FiniCB(Builder.saveIP());
 
-    // The continuation block is where code generation continues.
-    Builder.SetInsertPoint(NonCancellationBlock, NonCancellationBlock->begin());
-  }
 
-  return Builder.saveIP();
-}
-
 IRBuilder<>::InsertPoint OpenMPIRBuilder::CreateParallel(
     const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB,
     PrivatizeCallbackTy PrivCB, FinalizeCallbackTy FiniCB, Value *IfCondition,
     Value *NumThreads, omp::ProcBindKind ProcBind, bool IsCancellable) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   Constant *SrcLocStr = getOrCreateSrcLocStr(Loc);

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, CreateBarrier) {
   EXPECT_FALSE(Barrier->getCalledFunction()->doesNotFreeMemory());
 
   EXPECT_EQ(cast<CallInst>(Barrier)->getArgOperand(1), GTID);
 
   Builder.CreateUnreachable();
   EXPECT_FALSE(verifyModule(*M));
 }
 
+TEST_F(OpenMPIRBuilderTest, CreateCancel) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+
+  BasicBlock *CBB = BasicBlock::Create(Ctx, "", F);
+  new UnreachableInst(Ctx, CBB);
+  auto FiniCB = [&](InsertPointTy IP) {
+    ASSERT_NE(IP.getBlock(), nullptr);
+    ASSERT_EQ(IP.getBlock()->end(), IP.getPoint());
+    BranchInst::Create(CBB, IP.getBlock());
+  };
+  OMPBuilder.pushFinalizationCB({FiniCB, OMPD_parallel, true});
+
+  IRBuilder<> Builder(BB);
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP()});
+  auto NewIP = OMPBuilder.CreateCancel(Loc, nullptr, OMPD_parallel);
+  Builder.restoreIP(NewIP);
+  EXPECT_FALSE(M->global_empty());
+  EXPECT_EQ(M->size(), 3U);
+  EXPECT_EQ(F->size(), 4U);
+  EXPECT_EQ(BB->size(), 4U);
+
+  CallInst *GTID = dyn_cast<CallInst>(&BB->front());
+  EXPECT_NE(GTID, nullptr);
+  EXPECT_EQ(GTID->getNumArgOperands(), 1U);
+  EXPECT_EQ(GTID->getCalledFunction()->getName(), "__kmpc_global_thread_num");
+  EXPECT_FALSE(GTID->getCalledFunction()->doesNotAccessMemory());
+  EXPECT_FALSE(GTID->getCalledFunction()->doesNotFreeMemory());
+
+  CallInst *Cancel = dyn_cast<CallInst>(GTID->getNextNode());
+  EXPECT_NE(Cancel, nullptr);
+  EXPECT_EQ(Cancel->getNumArgOperands(), 3U);
+  EXPECT_EQ(Cancel->getCalledFunction()->getName(), "__kmpc_cancel");
+  EXPECT_FALSE(Cancel->getCalledFunction()->doesNotAccessMemory());
+  EXPECT_FALSE(Cancel->getCalledFunction()->doesNotFreeMemory());
+  EXPECT_EQ(Cancel->getNumUses(), 1U);
+  Instruction *CancelBBTI = Cancel->getParent()->getTerminator();
+  EXPECT_EQ(CancelBBTI->getNumSuccessors(), 2U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(0), NewIP.getBlock());
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->size(), 1U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->getTerminator()->getNumSuccessors(),
+            1U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->getTerminator()->getSuccessor(0),
+            CBB);
+
+  EXPECT_EQ(cast<CallInst>(Cancel)->getArgOperand(1), GTID);
+
+  OMPBuilder.popFinalizationCB();
+
+  Builder.CreateUnreachable();
+  EXPECT_FALSE(verifyModule(*M));
+}
+
+TEST_F(OpenMPIRBuilderTest, CreateCancelIfCond) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+
+  BasicBlock *CBB = BasicBlock::Create(Ctx, "", F);
+  new UnreachableInst(Ctx, CBB);
+  auto FiniCB = [&](InsertPointTy IP) {
+    ASSERT_NE(IP.getBlock(), nullptr);
+    ASSERT_EQ(IP.getBlock()->end(), IP.getPoint());
+    BranchInst::Create(CBB, IP.getBlock());
+  };
+  OMPBuilder.pushFinalizationCB({FiniCB, OMPD_parallel, true});
+
+  IRBuilder<> Builder(BB);
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP()});
+  auto NewIP = OMPBuilder.CreateCancel(Loc, Builder.getTrue(), OMPD_parallel);
+  Builder.restoreIP(NewIP);
+  EXPECT_FALSE(M->global_empty());
+  EXPECT_EQ(M->size(), 3U);
+  EXPECT_EQ(F->size(), 7U);
+  EXPECT_EQ(BB->size(), 1U);
+  ASSERT_TRUE(isa<BranchInst>(BB->getTerminator()));
+  ASSERT_EQ(BB->getTerminator()->getNumSuccessors(), 2U);
+  BB = BB->getTerminator()->getSuccessor(0);
+  EXPECT_EQ(BB->size(), 4U);
+
+
+  CallInst *GTID = dyn_cast<CallInst>(&BB->front());
+  EXPECT_NE(GTID, nullptr);
+  EXPECT_EQ(GTID->getNumArgOperands(), 1U);
+  EXPECT_EQ(GTID->getCalledFunction()->getName(), "__kmpc_global_thread_num");
+  EXPECT_FALSE(GTID->getCalledFunction()->doesNotAccessMemory());
+  EXPECT_FALSE(GTID->getCalledFunction()->doesNotFreeMemory());
+
+  CallInst *Cancel = dyn_cast<CallInst>(GTID->getNextNode());
+  EXPECT_NE(Cancel, nullptr);
+  EXPECT_EQ(Cancel->getNumArgOperands(), 3U);
+  EXPECT_EQ(Cancel->getCalledFunction()->getName(), "__kmpc_cancel");
+  EXPECT_FALSE(Cancel->getCalledFunction()->doesNotAccessMemory());
+  EXPECT_FALSE(Cancel->getCalledFunction()->doesNotFreeMemory());
+  EXPECT_EQ(Cancel->getNumUses(), 1U);
+  Instruction *CancelBBTI = Cancel->getParent()->getTerminator();
+  EXPECT_EQ(CancelBBTI->getNumSuccessors(), 2U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(0)->size(), 1U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(0)->getUniqueSuccessor(), NewIP.getBlock());
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->size(), 1U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->getTerminator()->getNumSuccessors(),
+            1U);
+  EXPECT_EQ(CancelBBTI->getSuccessor(1)->getTerminator()->getSuccessor(0),
+            CBB);
+
+  EXPECT_EQ(cast<CallInst>(Cancel)->getArgOperand(1), GTID);
+
+  OMPBuilder.popFinalizationCB();
+
+  Builder.CreateUnreachable();
+  EXPECT_FALSE(verifyModule(*M));
+}
+
 TEST_F(OpenMPIRBuilderTest, CreateCancelBarrier) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
 
   BasicBlock *CBB = BasicBlock::Create(Ctx, "", F);
   new UnreachableInst(Ctx, CBB);
   auto FiniCB = [&](InsertPointTy IP) {