[polly] [Polly] Consider the preconditions of the block of a user assumption (PR #226473)

[Timur Baidusenov via llvm-commits]

https://github.com/bai-tim updated https://github.com/llvm/llvm-project/pull/226473

>From 18252cc74693ddabadd97fbbf686f858ff334e75 Mon Sep 17 00:00:00 2001
From: Timur Baidusenov <timurbaidusenov at gmail.com>
Date: Fri, 25 Sep 2026 14:13:09 +0300
Subject: [PATCH 1/2] [Polly] Consider the preconditions of the block of a user
 assumption

An assumption inside the SCoP only holds where its block is executed, so its condition is built from the domain of the block, which may itself rely on preconditions, e.g. a condition such as (signed char)n < 10 that leads to the assumption is only modeled correctly if the truncation does not overflow. addUserAssumptions only checked the preconditions of the assumed expression, so such an assumption was added to the context, which gists the invalid context and removed the runtime check for these preconditions. Also add the assumption to the defined behavior context instead if the invalid domain of its block is not empty. Fixes #226419.

Assisted-by: Claude (Anthropic)
---
 polly/lib/Analysis/ScopBuilder.cpp            |  5 ++
 .../user-assumption-in-conditional-block.ll   | 60 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 polly/test/ScopInfo/user-assumption-in-conditional-block.ll

diff --git a/polly/lib/Analysis/ScopBuilder.cpp b/polly/lib/Analysis/ScopBuilder.cpp
index 7a52ee3481817..36ce6a82d7ee5 100644
--- a/polly/lib/Analysis/ScopBuilder.cpp
+++ b/polly/lib/Analysis/ScopBuilder.cpp
@@ -1464,6 +1464,11 @@ void ScopBuilder::addUserAssumptions(
 
     isl::set AssumptionInvalidDomain = AssumptionInvalidDomainMap[BB];
     bool HasPreconditions = !AssumptionInvalidDomain.is_empty();
+    // An assumption inside the SCoP only holds where its block is executed,
+    // i.e. it also relies on the preconditions of the domain of the block, such
+    // as those of the conditions that lead to it.
+    if (InScop && !BBInvalidDomain.is_empty())
+      HasPreconditions = true;
     InvalidDomainMap[BB] = BBInvalidDomain.unite(AssumptionInvalidDomain);
 
     if (!Valid)
diff --git a/polly/test/ScopInfo/user-assumption-in-conditional-block.ll b/polly/test/ScopInfo/user-assumption-in-conditional-block.ll
new file mode 100644
index 0000000000000..3cef4c3b5c11b
--- /dev/null
+++ b/polly/test/ScopInfo/user-assumption-in-conditional-block.ll
@@ -0,0 +1,60 @@
+; RUN: opt %loadNPMPolly '-passes=polly-custom<scops>' -polly-print-scops -disable-output < %s 2>&1 | FileCheck %s
+; RUN: opt %loadNPMPolly '-passes=polly-custom<ast>' -polly-print-ast -disable-output < %s 2>&1 | FileCheck %s --check-prefix=AST
+;
+; https://github.com/llvm/llvm-project/issues/226419
+;
+; The user assumption itself has no preconditions, but it is located in a
+; block whose domain is derived from (trunc i64 %n to i8), which is only
+; modeled correctly if the truncation does not overflow. Adding the assumption
+; to the context therefore relies on the 'n <= -129 or n >= 128' runtime
+; check, but the context is used to gist the invalid context, which removes
+; 'n <= -129' from it.
+; For n = -200, (signed char)n == 56, hence neither the assumption nor the
+; store is executed. The optimized code must not execute Stmt_then.
+;
+; void f(int *A, long n) {
+;   for (long i = 0; i < 4; i++)
+;     if ((signed char)n < 10) {
+;       __builtin_assume(n >= 0);
+;       A[i] = 1;
+;     }
+; }
+
+; CHECK:      Context:
+; CHECK-NEXT:   [n] -> {  : -9223372036854775808 <= n <= 9223372036854775807 }
+; CHECK:      Invalid Context:
+; CHECK-NEXT:   [n] -> {  : n <= -129 or n >= 128 }
+; CHECK:      Defined Behavior Context:
+; CHECK-NEXT:   [n] -> {  : 0 <= n <= 127 }
+; CHECK:      Domain :=
+; CHECK-NEXT:   [n] -> { Stmt_then[i0] : n <= 9 and 0 <= i0 <= 3 };
+
+; AST: if (1 && 0 == (n <= -129 || n >= 128))
+
+define void @f(ptr %A, i64 %n) {
+entry:
+  br label %for
+
+for:
+  %i = phi i64 [ 0, %entry ], [ %i.next, %latch ]
+  %t = trunc i64 %n to i8
+  %small = icmp slt i8 %t, 10
+  br i1 %small, label %then, label %latch
+
+then:
+  %c = icmp sge i64 %n, 0
+  call void @llvm.assume(i1 %c)
+  %gep = getelementptr inbounds i32, ptr %A, i64 %i
+  store i32 1, ptr %gep
+  br label %latch
+
+latch:
+  %i.next = add nuw nsw i64 %i, 1
+  %cmp = icmp slt i64 %i.next, 4
+  br i1 %cmp, label %for, label %exit
+
+exit:
+  ret void
+}
+
+declare void @llvm.assume(i1)

>From 115bb568740dfac8bf592df1644ad986a11a1740 Mon Sep 17 00:00:00 2001
From: Timur Baidusenov <timurbaidusenov at gmail.com>
Date: Sun, 4 Oct 2026 00:05:37 +0300
Subject: [PATCH 2/2] Compute HasPreconditions in one expression

---
 polly/lib/Analysis/ScopBuilder.cpp | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/polly/lib/Analysis/ScopBuilder.cpp b/polly/lib/Analysis/ScopBuilder.cpp
index 36ce6a82d7ee5..a9033e376875a 100644
--- a/polly/lib/Analysis/ScopBuilder.cpp
+++ b/polly/lib/Analysis/ScopBuilder.cpp
@@ -1463,12 +1463,8 @@ void ScopBuilder::addUserAssumptions(
     isl_set_free(Dom);
 
     isl::set AssumptionInvalidDomain = AssumptionInvalidDomainMap[BB];
-    bool HasPreconditions = !AssumptionInvalidDomain.is_empty();
-    // An assumption inside the SCoP only holds where its block is executed,
-    // i.e. it also relies on the preconditions of the domain of the block, such
-    // as those of the conditions that lead to it.
-    if (InScop && !BBInvalidDomain.is_empty())
-      HasPreconditions = true;
+    bool HasPreconditions = !AssumptionInvalidDomain.is_empty() ||
+                            (InScop && !BBInvalidDomain.is_empty());
     InvalidDomainMap[BB] = BBInvalidDomain.unite(AssumptionInvalidDomain);
 
     if (!Valid)
@@ -1503,8 +1499,11 @@ void ScopBuilder::addUserAssumptions(
 
     // scop->setContext is used to gist AssumedContext and InvalidContext. Both
     // add RTCs, so using setContext would remove the RTC that would ensure the
-    // correctness of AssumptionCtx. Using DefinedBehaviorContext which does not
-    // gist the other contexts.
+    // correctness of AssumptionCtx. Such an RTC exists if the translation of
+    // the assumption has preconditions, or if the assumption is inside the SCoP
+    // and the domain of its block, from which AssumptionCtx is built, has
+    // preconditions. Using DefinedBehaviorContext which does not gist the other
+    // contexts.
     // TODO: Use recordAssumption() for adding context/assumptions
     if (!HasPreconditions) {
       isl::set newContext =