From 453ef4e376a13f5ab8d9e2c6019a57101d8409a4 Mon Sep 17 00:00:00 2001
From: Amara Emerson <aemerson@apple.com>
Date: Mon, 2 Sep 2019 08:18:55 +0000
Subject: [PATCH] [AArch64][GlobalISel] Fix zext narrowScalar to use the right
 type when creating the merges.

Fixes PR43171.

llvm-svn: 370627
---
 llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp    |  8 +++++---
 .../AArch64/GlobalISel/legalize-sext-zext-128.mir  | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index c37b115..1fc10a2 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -620,13 +620,15 @@ LegalizerHelper::LegalizeResult LegalizerHelper::narrowScalar(MachineInstr &MI,
     if (TypeIdx != 0)
       return UnableToLegalize;
 
-    if (SizeOp0 % NarrowTy.getSizeInBits() != 0)
+    LLT SrcTy = MRI.getType(MI.getOperand(1).getReg());
+    uint64_t SizeOp1 = SrcTy.getSizeInBits();
+    if (SizeOp0 % SizeOp1 != 0)
       return UnableToLegalize;
 
     // Generate a merge where the bottom bits are taken from the source, and
     // zero everything else.
-    Register ZeroReg = MIRBuilder.buildConstant(NarrowTy, 0).getReg(0);
-    unsigned NumParts = SizeOp0 / NarrowTy.getSizeInBits();
+    Register ZeroReg = MIRBuilder.buildConstant(SrcTy, 0).getReg(0);
+    unsigned NumParts = SizeOp0 / SizeOp1;
     SmallVector<Register, 4> Srcs = {MI.getOperand(1).getReg()};
     for (unsigned Part = 1; Part < NumParts; ++Part)
       Srcs.push_back(ZeroReg);
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext-zext-128.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext-zext-128.mir
index d100167..6d4c969 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext-zext-128.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-sext-zext-128.mir
@@ -46,6 +46,28 @@ body:             |
 
 ...
 ---
+name:            narrow_zext_s128_from_s32
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $w0, $x1
+
+    ; CHECK-LABEL: name: narrow_zext_s128_from_s32
+    ; CHECK: liveins: $w0, $x1
+    ; CHECK: [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+    ; CHECK: [[COPY1:%[0-9]+]]:_(p0) = COPY $x1
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: [[MV:%[0-9]+]]:_(s128) = G_MERGE_VALUES [[COPY]](s32), [[C]](s32), [[C]](s32), [[C]](s32)
+    ; CHECK: G_STORE [[MV]](s128), [[COPY1]](p0) :: (store 16)
+    ; CHECK: RET_ReallyLR
+    %0:_(s32) = COPY $w0
+    %1:_(p0) = COPY $x1
+    %2:_(s128) = G_ZEXT %0(s32)
+    G_STORE %2(s128), %1(p0) :: (store 16)
+    RET_ReallyLR
+
+...
+---
 name:            narrow_zext_s192
 tracksRegLiveness: true
 body:             |
-- 
2.7.4