[JumpThreading] Clone noalias.scope.decl when threading blocks

When cloning instructions during jump threading, also clone and
adapt any declared scopes. This is primarily important when
threading loop exits, because we'll end up with two dominating
scope declarations in that case (at least after additional loop
rotation). This addresses a loose thread from
https://reviews.llvm.org/rG2556b413a7b8#975012.

Differential Revision: https://reviews.llvm.org/D97154

(cherry picked from commit 5e7e499b912d2c9ebaa91b5783ca123dbedeabcc)

4 files changed, 88 insertions, 0 deletions

diff --git a/llvm/include/llvm/Transforms/Utils/Cloning.h b/llvm/include/llvm/Transforms/Utils/Cloning.h
index 56aaa5d48e2a..aa960c625630 100644
--- a/llvm/include/llvm/Transforms/Utils/Cloning.h
+++ b/llvm/include/llvm/Transforms/Utils/Cloning.h
@@ -274,6 +274,13 @@ void updateProfileCallee(
 void identifyNoAliasScopesToClone(
     ArrayRef<BasicBlock *> BBs, SmallVectorImpl<MDNode *> &NoAliasDeclScopes);
 
+/// Find the 'llvm.experimental.noalias.scope.decl' intrinsics in the specified
+/// instruction range and extract their scope. These are candidates for
+/// duplication when cloning.
+void identifyNoAliasScopesToClone(
+    BasicBlock::iterator Start, BasicBlock::iterator End,
+    SmallVectorImpl<MDNode *> &NoAliasDeclScopes);
+
 /// Duplicate the specified list of noalias decl scopes.
 /// The 'Ext' string is added as an extension to the name.
 /// Afterwards, the ClonedScopes contains the mapping of the original scope
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 96aef90c1c1a..10b08b4e2224 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -2076,6 +2076,15 @@ JumpThreadingPass::cloneInstructions(BasicBlock::iterator BI,
     ValueMapping[PN] = NewPN;
   }
 
+  // Clone noalias scope declarations in the threaded block. When threading a
+  // loop exit, we would otherwise end up with two idential scope declarations
+  // visible at the same time.
+  SmallVector<MDNode *> NoAliasScopes;
+  DenseMap<MDNode *, MDNode *> ClonedScopes;
+  LLVMContext &Context = PredBB->getContext();
+  identifyNoAliasScopesToClone(BI, BE, NoAliasScopes);
+  cloneNoAliasScopes(NoAliasScopes, ClonedScopes, "thread", Context);
+
   // Clone the non-phi instructions of the source basic block into NewBB,
   // keeping track of the mapping and using it to remap operands in the cloned
   // instructions.
@@ -2084,6 +2093,7 @@ JumpThreadingPass::cloneInstructions(BasicBlock::iterator BI,
     New->setName(BI->getName());
     NewBB->getInstList().push_back(New);
     ValueMapping[&*BI] = New;
+    adaptNoAliasScopes(New, ClonedScopes, Context);
 
     // Remap operands to patch up intra-block references.
     for (unsigned i = 0, e = New->getNumOperands(); i != e; ++i)
diff --git a/llvm/lib/Transforms/Utils/CloneFunction.cpp b/llvm/lib/Transforms/Utils/CloneFunction.cpp
index 51a49574e55d..6ab061510a60 100644
--- a/llvm/lib/Transforms/Utils/CloneFunction.cpp
+++ b/llvm/lib/Transforms/Utils/CloneFunction.cpp
@@ -989,3 +989,11 @@ void llvm::identifyNoAliasScopesToClone(
       if (auto *Decl = dyn_cast<NoAliasScopeDeclInst>(&I))
         NoAliasDeclScopes.push_back(Decl->getScopeList());
 }
+
+void llvm::identifyNoAliasScopesToClone(
+    BasicBlock::iterator Start, BasicBlock::iterator End,
+    SmallVectorImpl<MDNode *> &NoAliasDeclScopes) {
+  for (Instruction &I : make_range(Start, End))
+    if (auto *Decl = dyn_cast<NoAliasScopeDeclInst>(&I))
+      NoAliasDeclScopes.push_back(Decl->getScopeList());
+}
diff --git a/llvm/test/Transforms/JumpThreading/noalias-scope-decl.ll b/llvm/test/Transforms/JumpThreading/noalias-scope-decl.ll
new file mode 100644
index 000000000000..b032afaaf313
--- /dev/null
+++ b/llvm/test/Transforms/JumpThreading/noalias-scope-decl.ll
@@ -0,0 +1,63 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -jump-threading < %s | FileCheck %s
+
+define void @test(i8* %ptr) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !0)
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[I_INC:%.*]], [[LATCH:%.*]] ]
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[I]], 100
+; CHECK-NEXT:    br i1 [[C]], label [[EXIT:%.*]], label [[LATCH]]
+; CHECK:       latch:
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !3)
+; CHECK-NEXT:    store i8 0, i8* [[PTR:%.*]], align 1, !noalias !0
+; CHECK-NEXT:    store i8 1, i8* [[PTR]], align 1, !noalias !3
+; CHECK-NEXT:    [[I_INC]] = add i32 [[I]], 1
+; CHECK-NEXT:    br label [[LOOP]]
+; CHECK:       exit:
+; CHECK-NEXT:    call void @llvm.experimental.noalias.scope.decl(metadata !5)
+; CHECK-NEXT:    store i8 0, i8* [[PTR]], align 1, !noalias !0
+; CHECK-NEXT:    store i8 1, i8* [[PTR]], align 1, !noalias !5
+; CHECK-NEXT:    ret void
+;
+entry:
+  call void @llvm.experimental.noalias.scope.decl(metadata !0)
+  br label %loop
+
+loop:
+  %i = phi i32 [ 0, %entry ], [ %i.inc, %latch ]
+  %c = icmp eq i32 %i, 100
+  br i1 %c, label %if, label %latch
+
+if:
+  br label %latch
+
+latch:
+  %p = phi i1 [ true, %if ], [ false, %loop ]
+  call void @llvm.experimental.noalias.scope.decl(metadata !3)
+  store i8 0, i8* %ptr, !noalias !0
+  store i8 1, i8* %ptr, !noalias !3
+  %i.inc = add i32 %i, 1
+  br i1 %p, label %exit, label %loop
+
+exit:
+  ret void
+}
+
+declare void @llvm.experimental.noalias.scope.decl(metadata)
+
+!0 = !{!1}
+!1 = distinct !{!1, !2, !"scope1"}
+!2 = distinct !{!2, !"domain"}
+!3 = !{!4}
+!4 = distinct !{!4, !2, !"scope2"}
+
+; CHECK: !0 = !{!1}
+; CHECK: !1 = distinct !{!1, !2, !"scope1"}
+; CHECK: !2 = distinct !{!2, !"domain"}
+; CHECK: !3 = !{!4}
+; CHECK: !4 = distinct !{!4, !2, !"scope2"}
+; CHECK: !5 = !{!6}
+; CHECK: !6 = distinct !{!6, !2, !"scope2:thread"}