From f52123b45406bbabd344d2f990ddadddc1898a00 Mon Sep 17 00:00:00 2001
From: Pete Cooper <peter_cooper@apple.com>
Date: Thu, 7 May 2015 19:21:36 +0000
Subject: [PATCH] [AArch64] Fix sext/zext folding in address arithmetic.

We were accidentally folding a sign/zero extend in to address arithmetic in a different BB when the extend wasn't available there.

Cross BB fast-isel isn't safe, so restrict this to only when the extend is in the same BB as the use.

llvm-svn: 236764
---
 llvm/lib/Target/AArch64/AArch64FastISel.cpp   | 61 ++++++++++---------
 .../AArch64/fast-isel-address-extends.ll      | 39 ++++++++++++
 2 files changed, 71 insertions(+), 29 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/fast-isel-address-extends.ll

diff --git a/llvm/lib/Target/AArch64/AArch64FastISel.cpp b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
index c9fbc85b6989..b80e2f88baef 100644
--- a/llvm/lib/Target/AArch64/AArch64FastISel.cpp
+++ b/llvm/lib/Target/AArch64/AArch64FastISel.cpp
@@ -664,20 +664,22 @@ bool AArch64FastISel::computeAddress(const Value *Obj, Address &Addr, Type *Ty)
     Addr.setExtendType(AArch64_AM::LSL);
 
     const Value *Src = U->getOperand(0);
-    if (const auto *I = dyn_cast<Instruction>(Src))
-      if (FuncInfo.MBBMap[I->getParent()] == FuncInfo.MBB)
-        Src = I;
-
-    // Fold the zext or sext when it won't become a noop.
-    if (const auto *ZE = dyn_cast<ZExtInst>(Src)) {
-      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>(Src)) {
-      if (!isIntExtFree(SE) && SE->getOperand(0)->getType()->isIntegerTy(32)) {
-        Addr.setExtendType(AArch64_AM::SXTW);
-        Src = SE->getOperand(0);
+    if (const auto *I = dyn_cast<Instruction>(Src)) {
+      if (FuncInfo.MBBMap[I->getParent()] == FuncInfo.MBB) {
+        // Fold the zext or sext when it won't become a noop.
+        if (const auto *ZE = dyn_cast<ZExtInst>(I)) {
+          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>(I)) {
+          if (!isIntExtFree(SE) &&
+              SE->getOperand(0)->getType()->isIntegerTy(32)) {
+            Addr.setExtendType(AArch64_AM::SXTW);
+            Src = SE->getOperand(0);
+          }
+        }
       }
     }
 
@@ -746,21 +748,22 @@ bool AArch64FastISel::computeAddress(const Value *Obj, Address &Addr, Type *Ty)
     Addr.setExtendType(AArch64_AM::LSL);
 
     const Value *Src = LHS;
-    if (const auto *I = dyn_cast<Instruction>(Src))
-      if (FuncInfo.MBBMap[I->getParent()] == FuncInfo.MBB)
-        Src = I;
-
-
-    // Fold the zext or sext when it won't become a noop.
-    if (const auto *ZE = dyn_cast<ZExtInst>(Src)) {
-      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>(Src)) {
-      if (!isIntExtFree(SE) && SE->getOperand(0)->getType()->isIntegerTy(32)) {
-        Addr.setExtendType(AArch64_AM::SXTW);
-        Src = SE->getOperand(0);
+    if (const auto *I = dyn_cast<Instruction>(Src)) {
+      if (FuncInfo.MBBMap[I->getParent()] == FuncInfo.MBB) {
+        // Fold the zext or sext when it won't become a noop.
+        if (const auto *ZE = dyn_cast<ZExtInst>(I)) {
+          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>(I)) {
+          if (!isIntExtFree(SE) &&
+              SE->getOperand(0)->getType()->isIntegerTy(32)) {
+            Addr.setExtendType(AArch64_AM::SXTW);
+            Src = SE->getOperand(0);
+          }
+        }
       }
     }
 
diff --git a/llvm/test/CodeGen/AArch64/fast-isel-address-extends.ll b/llvm/test/CodeGen/AArch64/fast-isel-address-extends.ll
new file mode 100644
index 000000000000..6a17ec502a02
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/fast-isel-address-extends.ll
@@ -0,0 +1,39 @@
+; RUN: llc %s -o - -O0 -verify-machineinstrs -fast-isel=true | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+target triple = "arm64-apple-ios8.0.0"
+
+; This test was trying to fold the sext %tmp142 in to the address arithmetic in %sunkaddr1.
+; This was incorrect as %.mux isn't available in the last bb.
+
+; CHECK: sxtw [[REG:x[0-9]+]]
+; CHECK: strh wzr, {{\[}}[[REG]], {{.*}}, lsl #1]
+
+; Function Attrs: nounwind optsize ssp
+define void @EdgeLoop(i32 %dir, i32 %edge, i32 %width, i16* %tmp89, i32 %tmp136, i16 %tmp144) #0 {
+bb:
+  %tmp2 = icmp eq i32 %dir, 0
+  %.mux = select i1 %tmp2, i32 %width, i32 1
+  %tmp142 = sext i32 %.mux to i64
+  %tmp151 = shl nsw i64 %tmp142, 1
+  %tmp153 = getelementptr inbounds i16, i16* %tmp89, i64 %tmp151
+  %tmp154 = load i16, i16* %tmp153, align 2
+  %tmp155 = zext i16 %tmp154 to i32
+  br i1 %tmp2, label %bb225, label %bb212
+
+bb212:                                            ; preds = %bb
+  store i16 %tmp144, i16* %tmp89, align 2
+  ret void
+
+bb225:                                            ; preds = %bb
+  %tmp248 = trunc i32 %tmp155 to i16
+  store i16 %tmp248, i16* %tmp89, align 2
+  %sunkaddr = ptrtoint i16* %tmp89 to i64
+  %sunkaddr1 = mul i64 %tmp142, 2
+  %sunkaddr2 = add i64 %sunkaddr, %sunkaddr1
+  %sunkaddr3 = inttoptr i64 %sunkaddr2 to i16*
+  store i16 0, i16* %sunkaddr3, align 2
+  ret void
+}
+
+attributes #0 = { nounwind optsize ssp }