[llvm] [llvm][LSR] Fix where invariant on ScaledReg & Scale is violated (PR #112576)

[Youngsuk Kim via llvm-commits]

https://github.com/JOE1994 updated https://github.com/llvm/llvm-project/pull/112576

>From cea80f598aad84ecee656fe0c7822eddf56fc822 Mon Sep 17 00:00:00 2001
From: Youngsuk Kim <youngsuk.kim at hpe.com>
Date: Wed, 16 Oct 2024 10:58:54 -0500
Subject: [PATCH 1/3] [llvm][LSR] Fix where invariant on ScaledReg & Scale is
 violated

Comments attached to the `ScaledReg` field of `struct Formula` explains that,
`ScaledReg` must be non-null when `Scale` is non-zero.

This fixes up a code path where this invariant is violated.
Also, add an assert to ensure this invariant holds true.

Without this patch, compiler aborts with the attached test case.

Fixes #76504
---
 .../Transforms/Scalar/LoopStrengthReduce.cpp  |  7 +++--
 .../Transforms/LoopStrengthReduce/pr76504.ll  | 29 +++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/LoopStrengthReduce/pr76504.ll

diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 575395eda1c5bb..28bd6318c0b9d7 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -622,6 +622,8 @@ static bool containsAddRecDependentOnLoop(const SCEV *S, const Loop &L) {
 /// representation.
 /// \see Formula::BaseRegs.
 bool Formula::isCanonical(const Loop &L) const {
+  assert((Scale == 0 || ScaledReg) && "ScaledReg must be non-null if Scale is non-zero");
+
   if (!ScaledReg)
     return BaseRegs.size() <= 1;
 
@@ -3973,9 +3975,10 @@ void LSRInstance::GenerateReassociationsImpl(LSRUse &LU, unsigned LUIdx,
       F.UnfoldedOffset =
           Immediate::getFixed((uint64_t)F.UnfoldedOffset.getFixedValue() +
                               InnerSumSC->getValue()->getZExtValue());
-      if (IsScaledReg)
+      if (IsScaledReg) {
         F.ScaledReg = nullptr;
-      else
+        F.Scale = 0;
+      } else
         F.BaseRegs.erase(F.BaseRegs.begin() + Idx);
     } else if (IsScaledReg)
       F.ScaledReg = InnerSum;
diff --git a/llvm/test/Transforms/LoopStrengthReduce/pr76504.ll b/llvm/test/Transforms/LoopStrengthReduce/pr76504.ll
new file mode 100644
index 00000000000000..ca1edc61a7e4de
--- /dev/null
+++ b/llvm/test/Transforms/LoopStrengthReduce/pr76504.ll
@@ -0,0 +1,29 @@
+; Reduced from https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=65323 by @RKSimon
+;
+; RUN: opt -S -passes=loop-reduce %s
+;
+; Make sure we don't trigger an assertion.
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+ at G = external global i32
+
+define void @foo() {
+bb8:
+  br label %bb30
+
+bb30:                                             ; preds = %bb30, %bb8
+  %l0 = phi i64 [ -2222, %bb8 ], [ %r23, %bb30 ]
+  %A22 = alloca i16, align 2
+  %r23 = add nuw i64 1, %l0
+  %G7 = getelementptr i16, ptr %A22, i64 %r23
+  %B15 = urem i64 %r23, %r23
+  %G6 = getelementptr i16, ptr %G7, i64 %B15
+  %B1 = urem i64 %r23, %r23
+  %B8 = sub i64 -1, %r23
+  %B18 = sub i64 %B8, %B1
+  %G5 = getelementptr i16, ptr %G6, i64 %B18
+  store ptr %G5, ptr undef, align 8
+  br label %bb30
+}

>From 8b34778707ab77ead506bcf765baa36a2d00a41f Mon Sep 17 00:00:00 2001
From: Youngsuk Kim <youngsuk.kim at hpe.com>
Date: Wed, 16 Oct 2024 11:31:45 -0500
Subject: [PATCH 2/3] apply clang-format

---
 llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
index 28bd6318c0b9d7..e55b8f6652e319 100644
--- a/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp
@@ -622,7 +622,8 @@ static bool containsAddRecDependentOnLoop(const SCEV *S, const Loop &L) {
 /// representation.
 /// \see Formula::BaseRegs.
 bool Formula::isCanonical(const Loop &L) const {
-  assert((Scale == 0 || ScaledReg) && "ScaledReg must be non-null if Scale is non-zero");
+  assert((Scale == 0 || ScaledReg) &&
+         "ScaledReg must be non-null if Scale is non-zero");
 
   if (!ScaledReg)
     return BaseRegs.size() <= 1;

>From eaa0576f091497571c325139ec91cb329a55e97d Mon Sep 17 00:00:00 2001
From: Youngsuk Kim <youngsuk.kim at hpe.com>
Date: Thu, 17 Oct 2024 09:23:48 -0500
Subject: [PATCH 3/3] Move test under X86/ & add FileCheck line

---
 llvm/test/Transforms/LoopStrengthReduce/{ => X86}/pr76504.ll | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)
 rename llvm/test/Transforms/LoopStrengthReduce/{ => X86}/pr76504.ll (92%)

diff --git a/llvm/test/Transforms/LoopStrengthReduce/pr76504.ll b/llvm/test/Transforms/LoopStrengthReduce/X86/pr76504.ll
similarity index 92%
rename from llvm/test/Transforms/LoopStrengthReduce/pr76504.ll
rename to llvm/test/Transforms/LoopStrengthReduce/X86/pr76504.ll
index ca1edc61a7e4de..94b9f7badb0f6d 100644
--- a/llvm/test/Transforms/LoopStrengthReduce/pr76504.ll
+++ b/llvm/test/Transforms/LoopStrengthReduce/X86/pr76504.ll
@@ -1,6 +1,6 @@
 ; Reduced from https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=65323 by @RKSimon
 ;
-; RUN: opt -S -passes=loop-reduce %s
+; RUN: opt -S -passes=loop-reduce %s | FileCheck %s
 ;
 ; Make sure we don't trigger an assertion.
 
@@ -10,6 +10,7 @@ target triple = "x86_64-unknown-linux-gnu"
 @G = external global i32
 
 define void @foo() {
+; CHECK-LABEL: foo
 bb8:
   br label %bb30