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

[Timur Baidusenov via llvm-commits]

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

A user assumption inside the SCoP only holds where its block is executed, so it also relies on the preconditions of the block's domain, e.g. (signed char)n < 10 in the condition leading to it is only modeled correctly if the truncation does not overflow. Such assumptions were added to the context, which removed the runtime check for these preconditions.
Soultion: they are added to defined behavior context if the invalid domain of their block is not empty.

Fixes #226419.

Assisted-by: Claude (Anthropic)

>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] [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)