From 2dd158d6553c252e6265f309f79274ecc9cae034 Mon Sep 17 00:00:00 2001
From: Johannes Doerfert <johannes@jdoerfert.de>
Date: Tue, 4 Oct 2022 07:45:54 -0700
Subject: [PATCH] [OpenMP] Make barrier elimination work in the presence of
 llvm.assume

Assumptions are droppable and eliminating them to eliminate barriers
seems reasonable.
---
 llvm/lib/Transforms/IPO/OpenMPOpt.cpp          | 16 +++++++++++++---
 llvm/test/Transforms/OpenMP/barrier_removal.ll |  3 +++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index f29ef16..7d1da92 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -1484,8 +1484,9 @@ private:
         // A barrier in a barrier pair is removeable if all instructions
         // between the barriers in the pair are side-effect free modulo the
         // barrier operation.
-        auto IsBarrierRemoveable = [&Kernel](BarrierInfo *StartBI,
-                                             BarrierInfo *EndBI) {
+        auto IsBarrierRemoveable = [&Kernel](
+                                       BarrierInfo *StartBI, BarrierInfo *EndBI,
+                                       SmallVector<AssumeInst *> &Assumptions) {
           assert(
               !StartBI->isImplicitExit() &&
               "Expected start barrier to be other than a kernel exit barrier");
@@ -1552,6 +1553,11 @@ private:
               continue;
             }
 
+            if (auto *AI = dyn_cast<AssumeInst>(I)) {
+              Assumptions.push_back(AI);
+              continue;
+            }
+
             if (auto *LI = dyn_cast<LoadInst>(I))
               if (LI->hasMetadata(LLVMContext::MD_invariant_load))
                 continue;
@@ -1577,12 +1583,16 @@ private:
           if (StartBI->isImplicit() && EndBI->isImplicit())
             continue;
 
-          if (!IsBarrierRemoveable(StartBI, EndBI))
+          SmallVector<AssumeInst *> Assumptions;
+          if (!IsBarrierRemoveable(StartBI, EndBI, Assumptions))
             continue;
 
           assert(!(StartBI->isImplicit() && EndBI->isImplicit()) &&
                  "Expected at least one explicit barrier to remove.");
 
+          for (auto *Assumption : Assumptions)
+            Assumption->eraseFromParent();
+
           // Remove an explicit barrier, check first, then second.
           if (!StartBI->isImplicit()) {
             LLVM_DEBUG(dbgs() << "Remove start barrier "
diff --git a/llvm/test/Transforms/OpenMP/barrier_removal.ll b/llvm/test/Transforms/OpenMP/barrier_removal.ll
index ebf8f54..4960966 100644
--- a/llvm/test/Transforms/OpenMP/barrier_removal.ll
+++ b/llvm/test/Transforms/OpenMP/barrier_removal.ll
@@ -9,6 +9,7 @@ declare i32 @llvm.nvvm.barrier0.and(i32)
 declare i32 @llvm.nvvm.barrier0.or(i32)
 declare i32 @llvm.nvvm.barrier0.popc(i32)
 declare void @llvm.amdgcn.s.barrier()
+declare void @llvm.assume(i1)
 
 ;.
 ; CHECK: @[[GC1:[a-zA-Z0-9_$"\\.-]+]] = constant i32 42
@@ -27,7 +28,9 @@ define void @pos_empty_1() {
 ; CHECK-LABEL: define {{[^@]+}}@pos_empty_1() {
 ; CHECK-NEXT:    ret void
 ;
+  call void @llvm.assume(i1 true)
   call void @unknown() "llvm.assume"="ompx_aligned_barrier"
+  call void @llvm.assume(i1 true)
   ret void
 }
 define void @pos_empty_2() {
-- 
2.7.4