From ef3722d8e989e1d25f7bab587a68bf2e41e9af72 Mon Sep 17 00:00:00 2001
From: Juergen Ributzka <juergen@apple.com>
Date: Tue, 7 Oct 2014 03:40:06 +0000
Subject: [PATCH] [FastISel][AArch64] Teach the address computation code to
 also fold sign-/zero-extends.

The code already folds sign-/zero-extends, but only if they are arguments to
mul and shift instructions. This extends the code to also fold them when they
are direct inputs.

llvm-svn: 219187
---
 llvm/lib/Target/AArch64/AArch64FastISel.cpp   | 29 ++++++++++++++++++
 .../test/CodeGen/AArch64/fast-isel-int-ext.ll | 30 +++++++------------
 2 files changed, 39 insertions(+), 20 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64FastISel.cpp b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
index 559a6bdf8ea8..110fd500e26f 100644
--- a/llvm/lib/Target/AArch64/AArch64FastISel.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
@@ -776,6 +776,35 @@ bool AArch64FastISel::computeAddress(const Value *Obj, Address &Addr, Type *Ty)
       }
     break;
   }
+  case Instruction::SExt:
+  case Instruction::ZExt: {
+    if (!Addr.getReg() || Addr.getOffsetReg())
+      break;
+
+    const Value *Src = nullptr;
+    // Fold the zext or sext when it won't become a noop.
+    if (const auto *ZE = dyn_cast<ZExtInst>(U)) {
+      if (!isIntExtFree(ZE) && ZE->getOperand(0)->getType()->isIntegerTy(32)) {
+        Addr.setExtendType(AArch64_AM::UXTW);
+        Src = ZE->getOperand(0);
+      }
+    } else if (const auto *SE = dyn_cast<SExtInst>(U)) {
+      if (!isIntExtFree(SE) && SE->getOperand(0)->getType()->isIntegerTy(32)) {
+        Addr.setExtendType(AArch64_AM::SXTW);
+        Src = SE->getOperand(0);
+      }
+    }
+
+    if (!Src)
+      break;
+
+    Addr.setShift(0);
+    unsigned Reg = getRegForValue(Src);
+    if (!Reg)
+      return false;
+    Addr.setOffsetReg(Reg);
+    return true;
+  }
   } // end switch
 
   if (Addr.getReg()) {
diff --git a/llvm/test/CodeGen/AArch64/fast-isel-int-ext.ll b/llvm/test/CodeGen/AArch64/fast-isel-int-ext.ll
index 4a783a87a70a..866febac2622 100644
--- a/llvm/test/CodeGen/AArch64/fast-isel-int-ext.ll
+++ b/llvm/test/CodeGen/AArch64/fast-isel-int-ext.ll
@@ -371,8 +371,7 @@ define i64 @load_register_sext_i32_to_i64(i64 %a, i64 %b) {
 ; Extend
 define i32 @load_extend_zext_i8_to_i32(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_zext_i8_to_i32
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrb w0, [x0, [[REG]]]
+; CHECK:       ldrb w0, [x0, w1, sxtw]
 ; CHECK-NOT:   uxtb
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -384,8 +383,7 @@ define i32 @load_extend_zext_i8_to_i32(i64 %a, i32 %b) {
 
 define i32 @load_extend_zext_i16_to_i32(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_zext_i16_to_i32
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrh w0, [x0, [[REG]]]
+; CHECK:       ldrh w0, [x0, w1, sxtw]
 ; CHECK-NOT:   uxth
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -397,8 +395,7 @@ define i32 @load_extend_zext_i16_to_i32(i64 %a, i32 %b) {
 
 define i64 @load_extend_zext_i8_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_zext_i8_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrb w0, [x0, [[REG]]]
+; CHECK:       ldrb w0, [x0, w1, sxtw]
 ; CHECK-NOT:   uxtb
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -410,8 +407,7 @@ define i64 @load_extend_zext_i8_to_i64(i64 %a, i32 %b) {
 
 define i64 @load_extend_zext_i16_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_zext_i16_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrh w0, [x0, [[REG]]]
+; CHECK:       ldrh w0, [x0, w1, sxtw]
 ; CHECK-NOT:   uxth
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -423,8 +419,7 @@ define i64 @load_extend_zext_i16_to_i64(i64 %a, i32 %b) {
 
 define i64 @load_extend_zext_i32_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_zext_i32_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldr w0, [x0, [[REG]]]
+; CHECK:       ldr w0, [x0, w1, sxtw]
 ; CHECK-NOT:   uxtw
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -436,8 +431,7 @@ define i64 @load_extend_zext_i32_to_i64(i64 %a, i32 %b) {
 
 define i32 @load_extend_sext_i8_to_i32(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_sext_i8_to_i32
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrsb w0, [x0, [[REG]]]
+; CHECK:       ldrsb w0, [x0, w1, sxtw]
 ; CHECK-NOT:   sxtb
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -449,8 +443,7 @@ define i32 @load_extend_sext_i8_to_i32(i64 %a, i32 %b) {
 
 define i32 @load_extend_sext_i16_to_i32(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_sext_i16_to_i32
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrsh w0, [x0, [[REG]]]
+; CHECK:       ldrsh w0, [x0, w1, sxtw]
 ; CHECK-NOT:   sxth
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -462,8 +455,7 @@ define i32 @load_extend_sext_i16_to_i32(i64 %a, i32 %b) {
 
 define i64 @load_extend_sext_i8_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_sext_i8_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrsb x0, [x0, [[REG]]]
+; CHECK:       ldrsb x0, [x0, w1, sxtw]
 ; CHECK-NOT:   sxtb
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -475,8 +467,7 @@ define i64 @load_extend_sext_i8_to_i64(i64 %a, i32 %b) {
 
 define i64 @load_extend_sext_i16_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_sext_i16_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrsh x0, [x0, [[REG]]]
+; CHECK:       ldrsh x0, [x0, w1, sxtw]
 ; CHECK-NOT:   sxth
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1
@@ -488,8 +479,7 @@ define i64 @load_extend_sext_i16_to_i64(i64 %a, i32 %b) {
 
 define i64 @load_extend_sext_i32_to_i64(i64 %a, i32 %b) {
 ; CHECK-LABEL: load_extend_sext_i32_to_i64
-; CHECK:       sxtw [[REG:x[0-9]+]], w1
-; CHECK-NEXT:  ldrsw x0, [x0, [[REG]]]
+; CHECK:       ldrsw x0, [x0, w1, sxtw]
 ; CHECK-NOT:   sxtw
   %1 = sext i32 %b to i64
   %2 = add i64 %a, %1