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

[via llvm-commits]

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

>From 93b460bf4f3ee81764903a756e92fa94726019a8 Mon Sep 17 00:00:00 2001
From: LiqinWeng <liqin.weng at spacemit.com>
Date: Tue, 28 Nov 2023 16:07:51 +0800
Subject: [PATCH] [InstCombine] Canonicalise SextADD + GEP

---
 .../InstCombine/InstructionCombining.cpp      | 22 +++++-
 llvm/test/Transforms/InstCombine/array.ll     | 74 +++++++++++++++++++
 2 files changed, 93 insertions(+), 3 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 e6088541349784b..146d6b6ca3f6e17 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2436,15 +2436,31 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
     if (match(GEP.getOperand(1),
               m_OneUse(m_Add(m_Value(Idx1), m_Value(Idx2))))) {
       //   %idx = add i64 %idx1, %idx2
-      //   %gep = getelementptr i32, i32* %ptr, i64 %idx
+      //   %gep = getelementptr i32, ptr %ptr, i64 %idx
       // as:
-      //   %newptr = getelementptr i32, i32* %ptr, i64 %idx1
-      //   %newgep = getelementptr i32, i32* %newptr, i64 %idx2
+      //   %newptr = getelementptr i32, ptr %ptr, i64 %idx1
+      //   %newgep = getelementptr i32, ptr %newptr, i64 %idx2
       auto *NewPtr = Builder.CreateGEP(GEP.getResultElementType(),
                                        GEP.getPointerOperand(), Idx1);
       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, ptr %ptr, i64 %sidx
+      // as:
+      // %newptr = getelementptr i32, ptr %ptr, i32 %idx1
+      // %newgep = getelementptr i32, ptr %newptr, i32 idx2
+      auto *NewPtr = Builder.CreateGEP(
+          GEP.getResultElementType(), GEP.getPointerOperand(),
+          Builder.CreateSExt(Idx1, GEP.getOperand(1)->getType()));
+      return GetElementPtrInst::Create(
+          GEP.getResultElementType(), NewPtr,
+          Builder.CreateSExt(C, GEP.getOperand(1)->getType()));
+    }
   }
 
   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..8bab3104fd8cd07
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/array.ll
@@ -0,0 +1,74 @@
+; 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 %ptr, i32 %a, i32 %b) {
+; CHECK-LABEL: define void @test(
+; CHECK-SAME: ptr [[PTR:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 [[A]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr i32, ptr [[PTR]], i64 [[TMP0]]
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i32, ptr [[TMP1]], i64 10
+; CHECK-NEXT:    store i32 [[B]], ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %add = add nsw i32 %a, 10
+  %idx = sext i32 %add to i64
+  %gep = getelementptr inbounds i32, ptr %ptr, i64 %idx
+  store i32 %b, ptr %gep
+  ret void
+}
+
+define  i32 @test_add_res_moreoneuse(ptr %ptr, i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @test_add_res_moreoneuse(
+; CHECK-SAME: ptr [[PTR:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[A]], 5
+; CHECK-NEXT:    [[IDX:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[PTR]], i64 [[IDX]]
+; CHECK-NEXT:    store i32 [[B]], ptr [[GEP]], align 4
+; CHECK-NEXT:    ret i32 [[ADD]]
+;
+entry:
+  %add = add nsw i32 %a, 5
+  %idx = sext i32 %add to i64
+  %gep = getelementptr inbounds i32, ptr %ptr, i64 %idx
+  store i32 %b, ptr %gep
+  ret i32 %add
+}
+
+define void @test_addop_nonsw_flag(ptr %ptr, i32 %a, i32 %b) {
+; CHECK-LABEL: define void @test_addop_nonsw_flag(
+; CHECK-SAME: ptr [[PTR:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[A]], 10
+; CHECK-NEXT:    [[IDX:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[PTR]], i64 [[IDX]]
+; CHECK-NEXT:    store i32 [[B]], ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %add = add i32 %a, 10
+  %idx = sext i32 %add to i64
+  %gep = getelementptr inbounds i32, ptr %ptr, i64 %idx
+  store i32 %b, ptr %gep
+  ret void
+}
+
+define void @test_add_op2_not_constant(ptr %ptr, i32 %a, i32 %b) {
+; CHECK-LABEL: define void @test_add_op2_not_constant(
+; CHECK-SAME: ptr [[PTR:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADD:%.*]] = add i32 [[A]], [[B]]
+; CHECK-NEXT:    [[IDX:%.*]] = sext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i32, ptr [[PTR]], i64 [[IDX]]
+; CHECK-NEXT:    store i32 [[B]], ptr [[GEP]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %add = add i32 %a, %b
+  %idx = sext i32 %add to i64
+  %gep = getelementptr inbounds i32, ptr %ptr, i64 %idx
+  store i32 %b, ptr %gep
+  ret void
+}