GlobalISel: legalize narrow G_SELECTS on AArch64.

Otherwise there aren't any patterns to select them.

llvm-svn: 294261

3 files changed, 42 insertions, 6 deletions

diff --git a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
index 27b6dc755a2..aeb89a41aa9 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp
@@ -343,6 +343,29 @@ LegalizerHelper::widenScalar(MachineInstr &MI, unsigned TypeIdx, LLT WideTy) {
     MI.eraseFromParent();
     return Legalized;
   }
+  case TargetOpcode::G_SELECT: {
+    if (TypeIdx != 0)
+      return UnableToLegalize;
+
+    // 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.
+    unsigned Src1Ext = MRI.createGenericVirtualRegister(WideTy);
+    unsigned Src2Ext = MRI.createGenericVirtualRegister(WideTy);
+    MIRBuilder.buildAnyExt(Src1Ext, MI.getOperand(2).getReg());
+    MIRBuilder.buildAnyExt(Src2Ext, MI.getOperand(3).getReg());
+
+    unsigned DstExt = MRI.createGenericVirtualRegister(WideTy);
+    MIRBuilder.buildInstr(TargetOpcode::G_SELECT)
+        .addDef(DstExt)
+        .addReg(MI.getOperand(1).getReg())
+        .addUse(Src1Ext)
+        .addUse(Src2Ext);
+
+    MIRBuilder.buildTrunc(MI.getOperand(0).getReg(), DstExt);
+    MI.eraseFromParent();
+    return Legalized;
+  }
   case TargetOpcode::G_FPTOSI:
   case TargetOpcode::G_FPTOUI: {
     if (TypeIdx != 0)
diff --git a/llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp b/llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp
index 838cb673b1a..3a97406c0d8 100644
--- a/llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp
@@ -182,7 +182,10 @@ AArch64LegalizerInfo::AArch64LegalizerInfo() {
   setAction({G_BRINDIRECT, p0}, Legal);
 
   // Select
-  for (auto Ty : {s1, s8, s16, s32, s64, p0})
+  for (auto Ty : {s1, s8, s16})
+    setAction({G_SELECT, Ty}, WidenScalar);
+
+  for (auto Ty : {s32, s64, p0})
     setAction({G_SELECT, Ty}, Legal);
 
   setAction({G_SELECT, 1, s1}, Legal);
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-simple.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-simple.mir
index 17685fa9bf1..f06e998453d 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalize-simple.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalize-simple.mir
@@ -52,14 +52,24 @@ body: |
 
   bb.1.next:
 
-    ; CHECK: %7(s1) = G_SELECT %1(s1), %1, %1
-    ; CHECK: %8(s8) = G_SELECT %1(s1), %2, %2
-    ; CHECK: %9(s16) = G_SELECT %1(s1), %3, %3
-    ; CHECK: %10(s32) = G_SELECT %1(s1), %4, %4
-    ; CHECK: %11(s64) = G_SELECT %1(s1), %0, %0
+    ; CHECK: [[LHS:%[0-9]+]](s32) = G_ANYEXT %1(s1)
+    ; CHECK: [[RHS:%[0-9]+]](s32) = G_ANYEXT %1(s1)
+    ; CHECK: [[RES:%[0-9]+]](s32) = G_SELECT %1(s1), [[LHS]], [[RHS]]
+    ; CHECK: %7(s1) = G_TRUNC [[RES]](s32)
     %7(s1) = G_SELECT %1, %1, %1
+
+    ; CHECK: [[LHS:%[0-9]+]](s32) = G_ANYEXT %2(s8)
+    ; CHECK: [[RHS:%[0-9]+]](s32) = G_ANYEXT %2(s8)
+    ; CHECK: [[RES:%[0-9]+]](s32) = G_SELECT %1(s1), [[LHS]], [[RHS]]
+    ; CHECK: %8(s8) = G_TRUNC [[RES]](s32)
     %8(s8) = G_SELECT %1, %2, %2
+
+    ; CHECK: [[LHS:%[0-9]+]](s32) = G_ANYEXT %3(s16)
+    ; CHECK: [[RHS:%[0-9]+]](s32) = G_ANYEXT %3(s16)
+    ; CHECK: [[RES:%[0-9]+]](s32) = G_SELECT %1(s1), [[LHS]], [[RHS]]
+    ; CHECK: %9(s16) = G_TRUNC [[RES]](s32)
     %9(s16) = G_SELECT %1, %3, %3
+
     %10(s32) = G_SELECT %1, %4, %4
     %11(s64) = G_SELECT %1, %0, %0