[llvm] [CodeGenPrepare] Replace deleted sext instr with the promoted value. (PR #71058)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw updated https://github.com/llvm/llvm-project/pull/71058

>From 0f117f963ed07e36e926f3db7964b2a1b2f2e28e Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 2 Nov 2023 20:30:26 +0800
Subject: [PATCH 1/3] [CodeGenPrepare] Replace deleted sext instr with the
 promoted value.

---
 llvm/lib/CodeGen/CodeGenPrepare.cpp           | 12 ++++++++++
 .../X86/codegen-prepare-addrmode-sext.ll      | 22 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 187820717b6fd5c..9018b87835b775f 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -2634,6 +2634,15 @@ struct ExtAddrMode : public TargetLowering::AddrMode {
   void print(raw_ostream &OS) const;
   void dump() const;
 
+  void replaceWith(Value *From, Value *To) {
+    if (BaseReg == From)
+      BaseReg = To;
+    if (ScaledReg == From)
+      ScaledReg = To;
+    if (OriginalValue == From)
+      OriginalValue = To;
+  }
+
   FieldName compare(const ExtAddrMode &other) {
     // First check that the types are the same on each field, as differing types
     // is something we can't cope with later on.
@@ -4933,6 +4942,9 @@ bool AddressingModeMatcher::matchOperationAddr(User *AddrInst, unsigned Opcode,
       TPT.rollback(LastKnownGood);
       return false;
     }
+
+    // SExt has been deleted. Make sure it is not referenced by the AddrMode.
+    AddrMode.replaceWith(Ext, PromotedOperand);
     return true;
   }
   }
diff --git a/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll b/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
index 6e95c91e7398932..a7f6f803d58fd4c 100644
--- a/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
+++ b/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
@@ -507,3 +507,25 @@ define i8 @oneArgPromotionBlockSExtZExt(i1 %arg1, ptr %base) {
   %res = load i8, ptr %arrayidx
   ret i8 %res
 }
+
+; Check that we replace the deleted sext with the promoted value.
+; CHECK-LABEL: define void @pr70938(
+; CHECK-SAME: ptr [[F:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 0, 0
+; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[ADD]], 2
+; CHECK-NEXT:    [[SUNKADDR1:%.*]] = getelementptr i8, ptr [[F]], i64 [[SUNKADDR]]
+; CHECK-NEXT:    store i8 0, ptr [[SUNKADDR1]], align 1
+; CHECK-NEXT:    ret void
+define void @pr70938(ptr %f) {
+entry:
+  %add = add nsw i32 0, 0
+  %idxprom3 = sext i32 %add to i64
+  %arrayidx4 = getelementptr [2 x [1 x [2 x i8]]], ptr %f, i64 0, i64 %idxprom3
+  %arrayidx8 = getelementptr [2 x i8], ptr %arrayidx4, i64 0, i64 %idxprom3
+  br label %if.end
+
+if.end:                                           ; preds = %entry
+  store i8 0, ptr %arrayidx8, align 1
+  ret void
+}

>From 77f8e43e05dd7538a3dd3cc171ff15bdd8203fc8 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Thu, 2 Nov 2023 23:35:55 +0800
Subject: [PATCH 2/3] fixup! [CodeGenPrepare] Replace deleted sext instr with
 the promoted value.

---
 llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll b/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
index a7f6f803d58fd4c..a3944ee2040b503 100644
--- a/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
+++ b/llvm/test/CodeGen/X86/codegen-prepare-addrmode-sext.ll
@@ -512,14 +512,15 @@ define i8 @oneArgPromotionBlockSExtZExt(i1 %arg1, ptr %base) {
 ; CHECK-LABEL: define void @pr70938(
 ; CHECK-SAME: ptr [[F:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 0, 0
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 0, 1
 ; CHECK-NEXT:    [[SUNKADDR:%.*]] = mul i64 [[ADD]], 2
 ; CHECK-NEXT:    [[SUNKADDR1:%.*]] = getelementptr i8, ptr [[F]], i64 [[SUNKADDR]]
-; CHECK-NEXT:    store i8 0, ptr [[SUNKADDR1]], align 1
+; CHECK-NEXT:    [[SUNKADDR2:%.*]] = getelementptr i8, ptr [[SUNKADDR1]], i64 1
+; CHECK-NEXT:    store i8 0, ptr [[SUNKADDR2]], align 1
 ; CHECK-NEXT:    ret void
 define void @pr70938(ptr %f) {
 entry:
-  %add = add nsw i32 0, 0
+  %add = add nsw i32 0, 1
   %idxprom3 = sext i32 %add to i64
   %arrayidx4 = getelementptr [2 x [1 x [2 x i8]]], ptr %f, i64 0, i64 %idxprom3
   %arrayidx8 = getelementptr [2 x i8], ptr %arrayidx4, i64 0, i64 %idxprom3

>From 8479b9d6535128377ad8abbc5ab45066f7c6d51a Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Sat, 4 Nov 2023 02:12:19 +0800
Subject: [PATCH 3/3] fixup! [CodeGenPrepare] Replace deleted sext instr with
 the promoted value.

Add comments for `replaceWith`.
---
 llvm/lib/CodeGen/CodeGenPrepare.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 9018b87835b775f..1eb610bb281a1e6 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -2634,6 +2634,9 @@ struct ExtAddrMode : public TargetLowering::AddrMode {
   void print(raw_ostream &OS) const;
   void dump() const;
 
+  // Replace From in ExtAddrMode with To.
+  // E.g., SExt insts may be promoted and deleted. We should replace them with
+  // the promoted values.
   void replaceWith(Value *From, Value *To) {
     if (BaseReg == From)
       BaseReg = To;