[AArch64] Optimize add/sub with immediate

[benshi001]

Summary

Optimize ([add|sub] r, imm) -> ([ADD|SUB] ([ADD|SUB] r, #imm0, lsl #12), imm1),
if imm == (imm0<<12)+imm1. and both imm0 and imm1 are non-zero 12-bit unsigned
integers.

Optimize ([add|sub] r, imm) -> ([SUB|ADD] ([SUB|ADD] r, #imm0, lsl #12), imm1),
if imm == -(imm0<<12)-imm1, and both imm0 and imm1 are non-zero 12-bit unsigned
integers.

[benshi001]

llvm/test/CodeGen/AArch64/addimm-mulimm.ll

56–58

I think this is still a win, since an extra register w9 is saved, although the total amount of instructions does not change.

[benshi001]

It seems the build failures do not related to my modification.

[dmgreen]

Similar to D109963, this may make things worse if the code is in a loop and the MOVi is loop invariant. That patch adds a new AArch64 peephole optimization pass, and this sounds like this change would fit in there nicely too.

[benshi001]

In D110476#3023610, @dmgreen wrote:

Similar to D109963, this may make things worse if the code is in a loop and the MOVi is loop invariant. That patch adds a new AArch64 peephole optimization pass, and this sounds like this change would fit in there nicely too.

Thanks. I will fit my code to that pass after the D109963 is merged.

[benshi001]

Re-implemented in D111034.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

 //===-- AArch64ISelDAGToDAG.cpp - A dag to dag inst selector for AArch64 --===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines an instruction selector for the AArch64 target.
 //
 //===----------------------------------------------------------------------===//
 
+#include "AArch64ExpandImm.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64TargetMachine.h"
 #include "MCTargetDesc/AArch64AddressingModes.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/CodeGen/SelectionDAGISel.h"
 #include "llvm/IR/Function.h" // To access function attributes.
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Intrinsics.h"

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 
 def : InstAlias<"negs $dst, $src", (SUBSWrs GPR32:$dst, WZR, GPR32:$src, 0), 3>;
 def : InstAlias<"negs $dst, $src", (SUBSXrs GPR64:$dst, XZR, GPR64:$src, 0), 3>;
 def : InstAlias<"negs $dst, $src$shift",
                 (SUBSWrs GPR32:$dst, WZR, GPR32:$src, arith_shift32:$shift), 2>;
 def : InstAlias<"negs $dst, $src$shift",
                 (SUBSXrs GPR64:$dst, XZR, GPR64:$src, arith_shift64:$shift), 2>;
 
+// Optimize ([add|sub] r, imm) ->
+//          ([ADD|SUB] ([ADD|SUB] r, #imm0, lsl #12), #imm1) if possible.
+def PTwoPartsAddImm : PatLeaf<(imm), [{
+  int64_t I = N->getSExtValue();
+  // The imm must be in the form of ((imm0 << 12) + imm1), in which both imm0
+  // and imm1 are non-zero 12-bit unsigned int.
+  if ((I & ~0xffffffll) != 0 || (I & 0xfff000) == 0 || (I & 0xfff) == 0)
+    return false;
+  // The imm has to be composed by at least two instructions.
+  SmallVector<AArch64_IMM::ImmInsnModel, 4> Insn;
+  const EVT VT = N->getValueType(0);
+  AArch64_IMM::expandMOVImm(N->getZExtValue(), VT.getSizeInBits(), Insn);
+  return Insn.size() > 1;
+}]>;
+
+def PTwoPartsAddImmHigh : SDNodeXForm<imm, [{
+  int64_t I = N->getSExtValue() >> 12;
+  return CurDAG->getTargetConstant(I & 0xfff, SDLoc(N), N->getValueType(0));
+}]>;
+
+def PTwoPartsAddImmLow : SDNodeXForm<imm, [{
+  int64_t I = N->getSExtValue();
+  return CurDAG->getTargetConstant(I & 0xfff, SDLoc(N), N->getValueType(0));
+}]>;
+
+def : Pat<(add GPR64:$R, PTwoPartsAddImm:$imm),
+          (ADDXri (ADDXri GPR64:$R, (PTwoPartsAddImmHigh PTwoPartsAddImm:$imm),
+                          12),
+                  (PTwoPartsAddImmLow PTwoPartsAddImm:$imm), 0)>;
+def : Pat<(add GPR32:$R, PTwoPartsAddImm:$imm),
+          (ADDWri (ADDWri GPR32:$R, (PTwoPartsAddImmHigh PTwoPartsAddImm:$imm),
+                          12),
+                  (PTwoPartsAddImmLow PTwoPartsAddImm:$imm), 0)>;
+def : Pat<(sub GPR64:$R, PTwoPartsAddImm:$imm),
+          (SUBXri (SUBXri GPR64:$R, (PTwoPartsAddImmHigh PTwoPartsAddImm:$imm),
+                          12),
+                  (PTwoPartsAddImmLow PTwoPartsAddImm:$imm), 0)>;
+def : Pat<(sub GPR32:$R, PTwoPartsAddImm:$imm),
+          (SUBWri (SUBWri GPR32:$R, (PTwoPartsAddImmHigh PTwoPartsAddImm:$imm),
+                          12),
+                  (PTwoPartsAddImmLow PTwoPartsAddImm:$imm), 0)>;
+
+// Optimize ([add|sub] r, imm) ->
+//          ([SUB|ADD] ([ADD|SUB] r, #imm0, lsl #12), #imm1) if possible.
+def NTwoPartsAddImm : PatLeaf<(imm), [{
+  int64_t I = -N->getSExtValue();
+  // The -imm must be in the form of ((imm0 << 12) + imm1), in which both imm0
+  // and imm1 are non-zero 12-bit unsigned int.
+  if ((I & ~0xffffffll) != 0 || (I & 0xfff000) == 0 || (I & 0xfff) == 0)
+    return false;
+  // The imm has to be composed by at least two instructions.
+  SmallVector<AArch64_IMM::ImmInsnModel, 4> Insn;
+  const EVT VT = N->getValueType(0);
+  AArch64_IMM::expandMOVImm(N->getZExtValue(), VT.getSizeInBits(), Insn);
+  return Insn.size() > 1;
+}]>;
+
+def NTwoPartsAddImmHigh : SDNodeXForm<imm, [{
+  int64_t I = (-N->getSExtValue()) >> 12;
+  return CurDAG->getTargetConstant(I & 0xfff, SDLoc(N), N->getValueType(0));
+}]>;
+
+def NTwoPartsAddImmLow : SDNodeXForm<imm, [{
+  int64_t I = -N->getSExtValue();
+  return CurDAG->getTargetConstant(I & 0xfff, SDLoc(N), N->getValueType(0));
+}]>;
+
+def : Pat<(add GPR64:$R, NTwoPartsAddImm:$imm),
+          (SUBXri (SUBXri GPR64:$R, (NTwoPartsAddImmHigh NTwoPartsAddImm:$imm),
+                          12),
+                  (NTwoPartsAddImmLow NTwoPartsAddImm:$imm), 0)>;
+def : Pat<(add GPR32:$R, NTwoPartsAddImm:$imm),
+          (SUBWri (SUBWri GPR32:$R, (NTwoPartsAddImmHigh NTwoPartsAddImm:$imm),
+                          12),
+                  (NTwoPartsAddImmLow NTwoPartsAddImm:$imm), 0)>;
+def : Pat<(sub GPR64:$R, NTwoPartsAddImm:$imm),
+          (ADDXri (ADDXri GPR64:$R, (NTwoPartsAddImmHigh NTwoPartsAddImm:$imm),
+                          12),
+                  (NTwoPartsAddImmLow NTwoPartsAddImm:$imm), 0)>;
+def : Pat<(sub GPR32:$R, NTwoPartsAddImm:$imm),
+          (ADDWri (ADDWri GPR32:$R, (NTwoPartsAddImmHigh NTwoPartsAddImm:$imm),
+                          12),
+                  (NTwoPartsAddImmLow NTwoPartsAddImm:$imm), 0)>;
 
 // Unsigned/Signed divide
 defm UDIV : Div<0, "udiv", udiv>;
 defm SDIV : Div<1, "sdiv", sdiv>;
 
 def : Pat<(int_aarch64_udiv GPR32:$Rn, GPR32:$Rm), (UDIVWr GPR32:$Rn, GPR32:$Rm)>;
 def : Pat<(int_aarch64_udiv GPR64:$Rn, GPR64:$Rm), (UDIVXr GPR64:$Rn, GPR64:$Rm)>;
 def : Pat<(int_aarch64_sdiv GPR32:$Rn, GPR32:$Rm), (SDIVWr GPR32:$Rn, GPR32:$Rm)>;

llvm/test/CodeGen/AArch64/addimm-mulimm.ll

; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, -31
   %tmp1 = mul i32 %tmp0, 37
   ret i32 %tmp1
 }
 
 define i64 @addimm_mulimm_accept_10(i64 %a) {
 ; CHECK-LABEL: addimm_mulimm_accept_10:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #32888
 ; CHECK-NEXT:    mov w8, #37
-; CHECK-NEXT:    movk w9, #17, lsl #16
-; CHECK-NEXT:    madd x0, x0, x8, x9
+; CHECK-NEXT:    mul x8, x0, x8
+; CHECK-NEXT:    add x8, x8, #280, lsl #12 // =1146880
+; CHECK-NEXT:    add x0, x8, #120
 ; CHECK-NEXT:    ret
   %tmp0 = add i64 %a, 31000
   %tmp1 = mul i64 %tmp0, 37
   ret i64 %tmp1
 }
 
 define i64 @addimm_mulimm_accept_11(i64 %a) {
 ; CHECK-LABEL: addimm_mulimm_accept_11:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x9, #-32888
 ; CHECK-NEXT:    mov w8, #37
-; CHECK-NEXT:    movk x9, #65518, lsl #16
-; CHECK-NEXT:    madd x0, x0, x8, x9
+; CHECK-NEXT:    mul x8, x0, x8
+; CHECK-NEXT:    sub x8, x8, #280, lsl #12 // =1146880
+; CHECK-NEXT:    sub x0, x8, #120
 ; CHECK-NEXT:    ret
   %tmp0 = add i64 %a, -31000
   %tmp1 = mul i64 %tmp0, 37
   ret i64 %tmp1
 }
 
 define signext i32 @addimm_mulimm_accept_12(i32 signext %a) {
 ; CHECK-LABEL: addimm_mulimm_accept_12:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #32888
 ; CHECK-NEXT:    mov w8, #37
-; CHECK-NEXT:    movk w9, #17, lsl #16
-; CHECK-NEXT:    madd w0, w0, w8, w9
+; CHECK-NEXT:    mul w8, w0, w8
+; CHECK-NEXT:    add w8, w8, #280, lsl #12 // =1146880
+; CHECK-NEXT:    add w0, w8, #120
 ; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, 31000
   %tmp1 = mul i32 %tmp0, 37
   ret i32 %tmp1
 }
 
 define signext i32 @addimm_mulimm_accept_13(i32 signext %a) {
 ; CHECK-LABEL: addimm_mulimm_accept_13:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w9, #32648
 ; CHECK-NEXT:    mov w8, #37
-; CHECK-NEXT:    movk w9, #65518, lsl #16
-; CHECK-NEXT:    madd w0, w0, w8, w9
+; CHECK-NEXT:    mul w8, w0, w8
+; CHECK-NEXT:    sub w8, w8, #280, lsl #12 // =1146880
+; CHECK-NEXT:    sub w0, w8, #120
 ; CHECK-NEXT:    ret
   %tmp0 = add i32 %a, -31000
   %tmp1 = mul i32 %tmp0, 37
   ret i32 %tmp1
 }
 
 define i64 @addimm_mulimm_reject_00(i64 %a) {
 ; CHECK-LABEL: addimm_mulimm_reject_00:

llvm/test/CodeGen/AArch64/addsub.ll

; CHECK-NEXT:    ret
   store i64 %newval64, i64* @var_i64
 
   ret void
 }
 
 define i64 @add_two_parts_imm_i64(i64 %a) {
 ; CHECK-LABEL: add_two_parts_imm_i64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #42325
-; CHECK-NEXT:    movk w8, #170, lsl #16
-; CHECK-NEXT:    add x0, x0, x8
+; CHECK-NEXT:    add x8, x0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    add x0, x8, #1365
 ; CHECK-NEXT:    ret
   %b = add i64 %a, 11183445
   ret i64 %b
 }
 
 define i32 @add_two_parts_imm_i32(i32 %a) {
 ; CHECK-LABEL: add_two_parts_imm_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #42325
-; CHECK-NEXT:    movk w8, #170, lsl #16
-; CHECK-NEXT:    add w0, w0, w8
+; CHECK-NEXT:    add w8, w0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    add w0, w8, #1365
 ; CHECK-NEXT:    ret
   %b = add i32 %a, 11183445
   ret i32 %b
 }
 
 define i64 @add_two_parts_imm_i64_neg(i64 %a) {
 ; CHECK-LABEL: add_two_parts_imm_i64_neg:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x8, #-42325
-; CHECK-NEXT:    movk x8, #65365, lsl #16
-; CHECK-NEXT:    add x0, x0, x8
+; CHECK-NEXT:    sub x8, x0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    sub x0, x8, #1365
 ; CHECK-NEXT:    ret
   %b = add i64 %a, -11183445
   ret i64 %b
 }
 
 define i32 @add_two_parts_imm_i32_neg(i32 %a) {
 ; CHECK-LABEL: add_two_parts_imm_i32_neg:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #23211
-; CHECK-NEXT:    movk w8, #65365, lsl #16
-; CHECK-NEXT:    add w0, w0, w8
+; CHECK-NEXT:    sub w8, w0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    sub w0, w8, #1365
 ; CHECK-NEXT:    ret
   %b = add i32 %a, -11183445
   ret i32 %b
 }
 
 define i64 @sub_two_parts_imm_i64(i64 %a) {
 ; CHECK-LABEL: sub_two_parts_imm_i64:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov x8, #-42325
-; CHECK-NEXT:    movk x8, #65365, lsl #16
-; CHECK-NEXT:    add x0, x0, x8
+; CHECK-NEXT:    sub x8, x0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    sub x0, x8, #1365
 ; CHECK-NEXT:    ret
   %b = sub i64 %a, 11183445
   ret i64 %b
 }
 
 define i32 @sub_two_parts_imm_i32(i32 %a) {
 ; CHECK-LABEL: sub_two_parts_imm_i32:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #23211
-; CHECK-NEXT:    movk w8, #65365, lsl #16
-; CHECK-NEXT:    add w0, w0, w8
+; CHECK-NEXT:    sub w8, w0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    sub w0, w8, #1365
 ; CHECK-NEXT:    ret
   %b = sub i32 %a, 11183445
   ret i32 %b
 }
 
 define i64 @sub_two_parts_imm_i64_neg(i64 %a) {
 ; CHECK-LABEL: sub_two_parts_imm_i64_neg:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #42325
-; CHECK-NEXT:    movk w8, #170, lsl #16
-; CHECK-NEXT:    add x0, x0, x8
+; CHECK-NEXT:    add x8, x0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    add x0, x8, #1365
 ; CHECK-NEXT:    ret
   %b = sub i64 %a, -11183445
   ret i64 %b
 }
 
 define i32 @sub_two_parts_imm_i32_neg(i32 %a) {
 ; CHECK-LABEL: sub_two_parts_imm_i32_neg:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    mov w8, #42325
-; CHECK-NEXT:    movk w8, #170, lsl #16
-; CHECK-NEXT:    add w0, w0, w8
+; CHECK-NEXT:    add w8, w0, #2730, lsl #12 // =11182080
+; CHECK-NEXT:    add w0, w8, #1365
 ; CHECK-NEXT:    ret
   %b = sub i32 %a, -11183445
   ret i32 %b
 }
 
 define void @testing() {
 ; CHECK-LABEL: testing:
 ; CHECK:       // %bb.0:

llvm/test/CodeGen/AArch64/srem-seteq-illegal-types.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
 
 define i1 @test_srem_odd(i29 %X) nounwind {
 ; CHECK-LABEL: test_srem_odd:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #33099
-; CHECK-NEXT:    mov w9, #24493
 ; CHECK-NEXT:    movk w8, #8026, lsl #16
-; CHECK-NEXT:    movk w9, #41, lsl #16
-; CHECK-NEXT:    madd w8, w0, w8, w9
+; CHECK-NEXT:    mul w8, w0, w8
+; CHECK-NEXT:    add w8, w8, #661, lsl #12 // =2707456
+; CHECK-NEXT:    add w8, w8, #4013
 ; CHECK-NEXT:    mov w9, #48987
 ; CHECK-NEXT:    and w8, w8, #0x1fffffff
 ; CHECK-NEXT:    movk w9, #82, lsl #16
 ; CHECK-NEXT:    cmp w8, w9
 ; CHECK-NEXT:    cset w0, lo
 ; CHECK-NEXT:    ret
   %srem = srem i29 %X, 99
   %cmp = icmp eq i29 %srem, 0

llvm/test/Transforms/CodeGenPrepare/AArch64/large-offset-gep.ll

 define void @test5([65536 x i32]** %s, i32 %n) {
 ; CHECK-LABEL: test5
 entry:
   %struct = load [65536 x i32]*, [65536 x i32]** %s
   br label %while_cond
 
 while_cond:
   %phi = phi i32 [ 0, %entry ], [ %i, %while_body ]
-; CHECK:     mov     w{{[0-9]+}}, #14464
+; CHECK-NOT: mov     w{{[0-9]+}}, #14464
 ; CHECK-NOT: mov     w{{[0-9]+}}, #14468
+; CHECK-NOT: movk    w{{[0-9]+}}, #1, lsl #16
+; CHECK:     add     x{{[0-9]+}}, x{{[0-9]+}}, #19, lsl #12
+; CHECK:     add     x{{[0-9]+}}, x{{[0-9]+}}, #2176
+; CHECK-NOT: add     x{{[0-9]+}}, x{{[0-9]+}}, #2180
   %gep0 = getelementptr [65536 x i32], [65536 x i32]* %struct, i64 0, i32 20000
   %gep1 = getelementptr [65536 x i32], [65536 x i32]* %struct, i64 0, i32 20001
   %cmp = icmp slt i32 %phi, %n
   br i1 %cmp, label %while_body, label %while_end
 
 while_body:
 ; CHECK:     str      w{{[0-9]+}}, [x{{[0-9]+}}, #4]
   %i = add i32 %phi, 1