[llvm] [DropUnnecessaryAssumes] Fix iterator invalidation. (PR #203765)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

<details>
<summary>Changes</summary>

registerAssumption() below can append to (and reallocate) the cache's assumption vector. Use integer index for indexing instead of using the iterator. Stop at the original count, so we don't reprocess assumes created during the loop.

---
Full diff: https://github.com/llvm/llvm-project/pull/203765.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp (+2-2) 
- (modified) llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll (+64) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
index 4a7144fe6c77a..7c3856856cd66 100644
--- a/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
+++ b/llvm/lib/Transforms/Scalar/DropUnnecessaryAssumes.cpp
@@ -67,8 +67,8 @@ DropUnnecessaryAssumesPass::run(Function &F, FunctionAnalysisManager &FAM) {
   AssumptionCache &AC = FAM.getResult<AssumptionAnalysis>(F);
   bool Changed = false;
 
-  for (const WeakVH &Elem : AC.assumptions()) {
-    auto *Assume = cast_or_null<AssumeInst>(Elem);
+    for (unsigned I = 0, E = AC.assumptions().size(); I != E; ++I) {
+    auto *Assume = cast_or_null<AssumeInst>(AC.assumptions()[I]);
     if (!Assume)
       continue;
 
diff --git a/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll b/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
index 43fa08c070828..079c0403227f3 100644
--- a/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
+++ b/llvm/test/Transforms/DropUnnecessaryAssumes/dereferenceable.ll
@@ -52,3 +52,67 @@ define i8 @test_dereferenceable_with_align_ptr_used(ptr %p, i64 %size) {
   %val = load i8, ptr %p
   ret i8 %val
 }
+
+; Make sure newly created assumes are handled properly.
+define i8 @test_dereferenceable_with_align_cache_realloc(ptr %p, ptr %q, i1 %c) {
+; CHECK-LABEL: define i8 @test_dereferenceable_with_align_cache_realloc(
+; CHECK-SAME: ptr [[P:%.*]], ptr [[Q:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[P]], align 1
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    ret i8 [[V]]
+;
+; DROP-DEREF-LABEL: define i8 @test_dereferenceable_with_align_cache_realloc(
+; DROP-DEREF-SAME: ptr [[P:%.*]], ptr [[Q:%.*]], i1 [[C:%.*]]) {
+; DROP-DEREF-NEXT:    [[V:%.*]] = load i8, ptr [[P]], align 1
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    call void @llvm.assume(i1 true) [ "align"(ptr [[P]], i64 8) ]
+; DROP-DEREF-NEXT:    ret i8 [[V]]
+;
+  %v = load i8, ptr %p
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 true) [ "dereferenceable"(ptr %q, i64 8), "align"(ptr %p, i64 8) ]
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  call void @llvm.assume(i1 %c)
+  ret i8 %v
+}

``````````

</details>


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