From e3ad0db135e354ab3b6b734834c81e5b9e4cb07d Mon Sep 17 00:00:00 2001
From: Justin Bogner <mail@justinbogner.com>
Date: Thu, 19 Jan 2017 07:51:17 +0000
Subject: [PATCH] GlobalISel: Implement widening for shifts

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@292476 91177308-0d34-0410-b5e6-96231b3b80d8
---
 lib/CodeGen/GlobalISel/LegalizerHelper.cpp    | 14 +++---
 .../AArch64/GlobalISel/legalize-shift.mir     | 44 +++++++++++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 test/CodeGen/AArch64/GlobalISel/legalize-shift.mir

diff --git a/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index 443c6e1ee40..0f23c2f0358 100644
--- a/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -222,7 +222,8 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
   case TargetOpcode::G_MUL:
   case TargetOpcode::G_OR:
   case TargetOpcode::G_XOR:
-  case TargetOpcode::G_SUB: {
+  case TargetOpcode::G_SUB:
+  case TargetOpcode::G_SHL: {
     // Perform operation at larger width (any extension is fine here, high bits
     // don't affect the result) and then truncate the result back to the
     // original type.
@@ -242,10 +243,13 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
     return Legalized;
   }
   case TargetOpcode::G_SDIV:
-  case TargetOpcode::G_UDIV: {
-    unsigned ExtOp = MI.getOpcode() == TargetOpcode::G_SDIV
-                          ? TargetOpcode::G_SEXT
-                          : TargetOpcode::G_ZEXT;
+  case TargetOpcode::G_UDIV:
+  case TargetOpcode::G_ASHR:
+  case TargetOpcode::G_LSHR: {
+    unsigned ExtOp = MI.getOpcode() == TargetOpcode::G_SDIV ||
+                             MI.getOpcode() == TargetOpcode::G_ASHR
+                         ? TargetOpcode::G_SEXT
+                         : TargetOpcode::G_ZEXT;
 
     unsigned LHSExt = MRI.createGenericVirtualRegister(WideTy);
     MIRBuilder.buildInstr(ExtOp).addDef(LHSExt).addUse(
diff --git a/test/CodeGen/AArch64/GlobalISel/legalize-shift.mir b/test/CodeGen/AArch64/GlobalISel/legalize-shift.mir
new file mode 100644
index 00000000000..673b23562ff
--- /dev/null
+++ b/test/CodeGen/AArch64/GlobalISel/legalize-shift.mir
@@ -0,0 +1,44 @@
+# RUN: llc -O0 -run-pass=legalizer -global-isel %s -o - 2>&1 | FileCheck %s
+
+--- |
+  target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64--"
+  define void @test_shift() {
+  entry:
+    ret void
+  }
+...
+
+---
+name:            test_shift
+registers:
+  - { id: 0, class: _ }
+  - { id: 1, class: _ }
+  - { id: 2, class: _ }
+  - { id: 3, class: _ }
+  - { id: 4, class: _ }
+  - { id: 5, class: _ }
+  - { id: 6, class: _ }
+body: |
+  bb.0.entry:
+    liveins: %x0, %x1, %x2, %x3
+    %0(s64) = COPY %x0
+    %1(s64) = COPY %x1
+    %2(s8) = G_TRUNC %0
+    %3(s8) = G_TRUNC %1
+
+    ; CHECK: [[LHS32:%[0-9]+]](s32) = G_SEXT %2
+    ; CHECK: [[RHS32:%[0-9]+]](s32) = G_SEXT %3
+    ; CHECK: [[RES32:%[0-9]+]](s32) = G_ASHR [[LHS32]], [[RHS32]]
+    ; CHECK: %4(s8) = G_TRUNC [[RES32]]
+    %4(s8) = G_ASHR %2, %3
+
+    ; CHECK: [[LHS32:%[0-9]+]](s32) = G_ZEXT %2
+    ; CHECK: [[RHS32:%[0-9]+]](s32) = G_ZEXT %3
+    ; CHECK: [[RES32:%[0-9]+]](s32) = G_LSHR [[LHS32]], [[RHS32]]
+    ; CHECK: %5(s8) = G_TRUNC [[RES32]]
+    %5(s8) = G_LSHR %2, %3
+
+    ; CHECK: %6(s8) = G_SHL %2, %3
+    %6(s8) = G_SHL %2, %3
+...
-- 
2.40.0