[llvm] [InstCombine] Canonicalise SextADD + GEP (PR #69581)

[via llvm-commits]

https://github.com/LiqinWeng updated https://github.com/llvm/llvm-project/pull/69581

>From 6d0d41b93b9b34af5914796ac0943194c8376bb0 Mon Sep 17 00:00:00 2001
From: LiqinWeng <liqin.weng at spacemit.com>
Date: Thu, 19 Oct 2023 17:12:44 +0800
Subject: [PATCH] [InstCombine] Canonicalise SextADD + GEP

---
 .../InstCombine/InstructionCombining.cpp      | 16 ++++++++
 llvm/test/Transforms/InstCombine/array.ll     | 38 +++++++++++++++++++
 polly/test/Support/dumpmodule.ll              | 19 +++++-----
 3 files changed, 63 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/Transforms/InstCombine/array.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 559eb2ef4795eb1..4574fb0c9cc375a 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2335,6 +2335,22 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
       return GetElementPtrInst::Create(GEP.getResultElementType(), NewPtr,
                                        Idx2);
     }
+
+    ConstantInt *C;
+    if (match(GEP.getOperand(1), m_OneUse(m_SExt(m_OneUse(m_NSWAdd(
+                                     m_Value(Idx1), m_ConstantInt(C))))))) {
+      // %add = add nsw i32 %idx1, idx2
+      // %sidx = sext i32 %add to i64
+      // %gep = getelementptr i32, i32* %ptr, i64 %sidx
+      // as :
+      // %newptr = getelementptr i32, i32* %ptr, i32 %idx1
+      // %newgep = getelementptr i32, i32* %newptr, i32 idx2
+      auto *NewPtr =
+          Builder.CreateGEP(GEP.getResultElementType(), GEP.getPointerOperand(),
+                            Idx1);
+      return GetElementPtrInst::Create(GEP.getResultElementType(), NewPtr,
+                                       C);
+    }
   }
 
   if (!GEP.isInBounds()) {
diff --git a/llvm/test/Transforms/InstCombine/array.ll b/llvm/test/Transforms/InstCombine/array.ll
new file mode 100644
index 000000000000000..bf768d1f8d9173c
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/array.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define void @test(ptr noundef %array2, i32 noundef signext %a, i32 noundef signext %b) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr noundef [[ARRAY2:%.*]], i32 noundef signext [[A:%.*]], i32 noundef signext [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[A]], 5
+; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[SUB:%.*]] = add nsw i32 [[A]], 4
+; CHECK-NEXT:    [[IDXPROM1:%.*]] = sext i32 [[SUB]] to i64
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds [50 x i32], ptr [[ARRAY2]], i64 [[IDXPROM]], i64 [[IDXPROM1]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    store i32 [[ADD3]], ptr [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i32 [[A]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr [50 x i32], ptr [[ARRAY2]], i64 [[TMP1]]
+; CHECK-NEXT:    [[ARRAYIDX8:%.*]] = getelementptr [50 x i32], ptr [[TMP2]], i64 25, i64 [[IDXPROM]]
+; CHECK-NEXT:    store i32 [[ADD]], ptr [[ARRAYIDX8]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %add = add nsw i32 %a, 5
+  %idxprom = sext i32 %add to i64
+  %arrayidx = getelementptr inbounds [50 x i32], ptr %array2, i64 %idxprom
+  %sub = sub nsw i32 %add, 1
+  %idxprom1 = sext i32 %sub to i64
+  %arrayidx2 = getelementptr inbounds [50 x i32], ptr %arrayidx, i64 0, i64 %idxprom1
+  %0 = load i32, ptr %arrayidx2, align 4
+  %add3 = add nsw i32 %0, 1
+  store i32 %add3, ptr %arrayidx2, align 4
+  %add4 = add nsw i32 %add, 20
+  %idxprom5 = sext i32 %add4 to i64
+  %arrayidx6 = getelementptr inbounds [50 x i32], ptr %array2, i64 %idxprom5
+  %arrayidx8 = getelementptr inbounds [50 x i32], ptr %arrayidx6, i64 0, i64 %idxprom
+  store i32 %add, ptr %arrayidx8, align 4
+  ret void
+}
diff --git a/polly/test/Support/dumpmodule.ll b/polly/test/Support/dumpmodule.ll
index 693fe4bc6cde3ea..38cb94b88095f0e 100644
--- a/polly/test/Support/dumpmodule.ll
+++ b/polly/test/Support/dumpmodule.ll
@@ -24,11 +24,11 @@ for:
   %j.cmp = icmp slt i32 %j, %n
   br i1 %j.cmp, label %body, label %exit
 
-    body:
-      %idx = add i32 %i, %j
-      %arrayidx = getelementptr inbounds double, ptr %A, i32 %idx
-      store double 42.0, ptr %arrayidx
-      br label %inc
+  body:
+  %idx = add i32 %i, %j
+  %arrayidx = getelementptr inbounds double, ptr %A, i32 %idx
+  store double 42.0, ptr %arrayidx
+  br label %inc
 
 inc:
   %j.inc = add nuw nsw i32 %j, 1
@@ -51,9 +51,9 @@ for:
   %i.cmp = icmp slt i32 %i, %n
   br i1 %i.cmp, label %body, label %exit
 
-    body:
-      call void @callee(i32 %n, ptr %A, i32 %i)
-      br label %inc
+  body:
+  call void @callee(i32 %n, ptr %A, i32 %i)
+  br label %inc
 
 inc:
   %j.inc = add nuw nsw i32 %i, 1
@@ -68,8 +68,7 @@ return:
 
 
 ; EARLY-LABEL: @callee(
-; AFTEREARLY:  polly.split_new_and_old:
-; EARLY:         store double 4.200000e+01, ptr %arrayidx
+; AFTEREARLY:  polly.stmt.body:
 ; EARLY-LABEL: @caller(
 ; EARLY:         call void @callee(