[llvm] [SCEV] Fix "quick and dirty" difference that could lead to assert (PR #70688)

[Danila Malyutin via llvm-commits]

https://github.com/danilaml created https://github.com/llvm/llvm-project/pull/70688

The old algorithm would remove all operands matching %step SCEV when it intended to only remove a single one. This lead to assert when SCEVAddExpr was of the form %step + %step and potential miscompiles in similar cases. Such SCEVs could be created when construction reached depth thresholds.

Fixes #70348

>From 14edb8f88f401fcef1d6d338bb7ce2e9014639c8 Mon Sep 17 00:00:00 2001
From: Danila Malyutin <dmalyutin at azul.com>
Date: Mon, 30 Oct 2023 19:35:56 +0300
Subject: [PATCH] [SCEV] Fix "quick and dirty" difference that could lead to
 assert

The old algorithm would remove all operands matching %step SCEV when it
intended to only remove a single one. This lead to assert when SCEVAddExpr
was of the form %step + %step and potential miscompiles in similar cases.
Such SCEVs could be created when construction reached depth thresholds.

Fixes #70348
---
 llvm/lib/Analysis/ScalarEvolution.cpp         | 13 +++++++----
 .../Transforms/LoopStrengthReduce/pr70348.ll  | 23 +++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopStrengthReduce/pr70348.ll

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index 2368003177e741c..f13e508a6c454bd 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -1335,11 +1335,14 @@ static const SCEV *getPreStartForExtend(const SCEVAddRecExpr *AR, Type *Ty,
 
   // Create an AddExpr for "PreStart" after subtracting Step. Full SCEV
   // subtraction is expensive. For this purpose, perform a quick and dirty
-  // difference, by checking for Step in the operand list.
-  SmallVector<const SCEV *, 4> DiffOps;
-  for (const SCEV *Op : SA->operands())
-    if (Op != Step)
-      DiffOps.push_back(Op);
+  // difference, by checking for Step in the operand list. Note, that
+  // SA might have repeated ops, like %a + %a + ..., so only remove one.
+  SmallVector<const SCEV *, 4> DiffOps(SA->operands());
+  for (auto It = DiffOps.begin(); It != DiffOps.end(); ++It)
+    if (*It == Step) {
+      DiffOps.erase(It);
+      break;
+    }
 
   if (DiffOps.size() == SA->getNumOperands())
     return nullptr;
diff --git a/llvm/test/Transforms/LoopStrengthReduce/pr70348.ll b/llvm/test/Transforms/LoopStrengthReduce/pr70348.ll
new file mode 100644
index 000000000000000..35b48a03a14e435
--- /dev/null
+++ b/llvm/test/Transforms/LoopStrengthReduce/pr70348.ll
@@ -0,0 +1,23 @@
+; RUN: opt -S -passes=loop-reduce -scalar-evolution-max-arith-depth=0 %s | FileCheck %s
+;
+; Make sure we don't trigger an assertion in SCEV here.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128-ni:1-p2:32:8:8:32-ni:2"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @test(i32 %phi) {
+; CHECK-LABEL: test
+bb:
+  br label %bb6
+
+bb6:                                              ; preds = %bb6, %bb
+  %phi7 = phi i32 [ 1, %bb ], [ %add44, %bb6 ]
+  %mul13 = mul i32 %phi7, %phi
+  %mul16 = mul i32 %mul13, 0
+  %add44 = add i32 %phi7, 1
+  br i1 true, label %bb51, label %bb6
+
+bb51:                                             ; preds = %bb6
+  unreachable
+}
+