[AArch64] Add patterns for zext/sext of shift amount.

Not sure this is the best fix, but it saves an instruction for certain
constructs involving variable shifts.

Differential Revision: https://reviews.llvm.org/D55572

llvm-svn: 351768

llvm/lib/Target/AArch64/AArch64InstrFormats.td

@@ -1817,6 +1817,14 @@ multiclass Shift<bits<2> shift_type, string asm, SDNode OpNode> {
 
   def : Pat<(i32 (OpNode GPR32:$Rn, (i64 (sext GPR32:$Rm)))),
             (!cast<Instruction>(NAME # "Wr") GPR32:$Rn, GPR32:$Rm)>;
+
+  def : Pat<(i64 (OpNode GPR64:$Rn, (i64 (sext GPR32:$Rm)))),
+            (!cast<Instruction>(NAME # "Xr") GPR64:$Rn,
+                (SUBREG_TO_REG (i32 0), GPR32:$Rm, sub_32))>;
+
+  def : Pat<(i64 (OpNode GPR64:$Rn, (i64 (zext GPR32:$Rm)))),
+            (!cast<Instruction>(NAME # "Xr") GPR64:$Rn,
+                (SUBREG_TO_REG (i32 0), GPR32:$Rm, sub_32))>;
 }
 
 class ShiftAlias<string asm, Instruction inst, RegisterClass regtype>

llvm/test/CodeGen/AArch64/shift-mod.ll

@@ -1,3 +1,4 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 < %s | FileCheck %s
 
 ; Check that we optimize out AND instructions and ADD/SUB instructions
@@ -6,19 +7,20 @@
 
 define i32 @test1(i32 %x, i64 %y) {
 ; CHECK-LABEL: test1:
-; CHECK-NOT: and
-; CHECK: lsr
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsr w0, w0, w1
+; CHECK-NEXT:    ret
   %sh_prom = trunc i64 %y to i32
   %shr = lshr i32 %x, %sh_prom
   ret i32 %shr
 }
 
 define i64 @test2(i32 %x, i64 %y) {
 ; CHECK-LABEL: test2:
-; CHECK-NOT: orr
-; CHECK-NOT: sub
-; CHECK: neg
-; CHECK: asr
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    neg w[[REG:[0-9]+]], w0
+; CHECK-NEXT:    asr x0, x1, x[[REG]]
+; CHECK-NEXT:    ret
   %sub9 = sub nsw i32 64, %x
   %sh_prom12.i = zext i32 %sub9 to i64
   %shr.i = ashr i64 %y, %sh_prom12.i
@@ -27,9 +29,46 @@ define i64 @test2(i32 %x, i64 %y) {
 
 define i64 @test3(i64 %x, i64 %y) {
 ; CHECK-LABEL: test3:
-; CHECK-NOT: add
-; CHECK: lsl
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    lsl x0, x1, x0
+; CHECK-NEXT:    ret
   %add = add nsw i64 64, %x
   %shl = shl i64 %y, %add
   ret i64 %shl
-}
+}
+
+define i64 @test4(i64 %y, i32 %s) {
+; CHECK-LABEL: test4:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    // kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    asr x0, x0, x1
+; CHECK-NEXT:    ret
+entry:
+  %sh_prom = zext i32 %s to i64
+  %shr = ashr i64 %y, %sh_prom
+  ret i64 %shr
+}
+
+define i64 @test5(i64 %y, i32 %s) {
+; CHECK-LABEL: test5:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    // kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    asr x0, x0, x1
+; CHECK-NEXT:    ret
+entry:
+  %sh_prom = sext i32 %s to i64
+  %shr = ashr i64 %y, %sh_prom
+  ret i64 %shr
+}
+
+define i64 @test6(i64 %y, i32 %s) {
+; CHECK-LABEL: test6:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    // kill: def $w1 killed $w1 def $x1
+; CHECK-NEXT:    lsl x0, x0, x1
+; CHECK-NEXT:    ret
+entry:
+  %sh_prom = sext i32 %s to i64
+  %shr = shl i64 %y, %sh_prom
+  ret i64 %shr
+}