[AArch64] Fix creation of invalid instructions with XZR register

[dmgreen]

Summary

A combination of GlobalISel and MachineCombiner can end up creating SUB xrz, (MOVI -2105098) instructions which have not been constant folded. The AArch64MIPeepholeOpt pass will then attempt to create ADD xzr, 513, lsl 12, which is not a valid instruction. This adds a bail out of the transform if the register is xzr/wzr.

[red1bluelost]

LGTM , interesting catch

[efriedma]

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

326

Is the transform assuming that the immediate is on the RHS (operand index 2)? If so, could you update the comment describing the transform to note that explicitly? Otherwise, it isn't really clear why operand index 1 is special.

[dmgreen]

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

326

Yep that could do with a better comment. Will do, thanks.

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

static bool splitAddSubImm(T Imm, unsigned RegSize, T &Imm0, T &Imm1) {
   return true;
 }
 
 template <typename T>
 bool AArch64MIPeepholeOpt::visitADDSUB(
     unsigned PosOpc, unsigned NegOpc, MachineInstr &MI) {
   // Try below transformation.
   //
-  // MOVi32imm + ADDWrr ==> ADDWri + ADDWri
-  // MOVi64imm + ADDXrr ==> ADDXri + ADDXri
+  // ADDWrr X, MOVi32imm ==> ADDWri + ADDWri
+  // ADDXrr X, MOVi64imm ==> ADDXri + ADDXri
   //
-  // MOVi32imm + SUBWrr ==> SUBWri + SUBWri
-  // MOVi64imm + SUBXrr ==> SUBXri + SUBXri
+  // SUBWrr X, MOVi32imm ==> SUBWri + SUBWri
+  // SUBXrr X, MOVi64imm ==> SUBXri + SUBXri
   //
   // The mov pseudo instruction could be expanded to multiple mov instructions
   // later. Let's try to split the constant operand of mov instruction into two
   // legal add/sub immediates. It makes only two ADD/SUB instructions intead of
   // multiple `mov` + `and/sub` instructions.
 
+  // We can sometimes have ADDWrr WZR, MULi32imm that have not been constant
+  // folded. Make sure that we don't generate invalid instructions that use XZR
+  // in those cases.
+  if (MI.getOperand(1).getReg() == AArch64::XZR ||
+      MI.getOperand(1).getReg() == AArch64::WZR)
+    return false;
+
   return splitTwoPartImm<T>(
       MI,
       [PosOpc, NegOpc](T Imm, unsigned RegSize, T &Imm0,
                        T &Imm1) -> std::optional<OpcodePair> {
         if (splitAddSubImm(Imm, RegSize, Imm0, Imm1))
           return std::make_pair(PosOpc, PosOpc);
         if (splitAddSubImm(-Imm, RegSize, Imm0, Imm1))
           return std::make_pair(NegOpc, NegOpc);
return splitTwoPartImm<T>(
       });
 }
 
 template <typename T>
 bool AArch64MIPeepholeOpt::visitADDSSUBS(
     OpcodePair PosOpcs, OpcodePair NegOpcs, MachineInstr &MI) {
   // Try the same transformation as ADDSUB but with additional requirement
   // that the condition code usages are only for Equal and Not Equal
+
+  if (MI.getOperand(1).getReg() == AArch64::XZR ||
+      MI.getOperand(1).getReg() == AArch64::WZR)
+    return false;
+
   return splitTwoPartImm<T>(
       MI,
       [PosOpcs, NegOpcs, &MI, &TRI = TRI,
        &MRI = MRI](T Imm, unsigned RegSize, T &Imm0,
                    T &Imm1) -> std::optional<OpcodePair> {
         OpcodePair OP;
         if (splitAddSubImm(Imm, RegSize, Imm0, Imm1))
           OP = PosOpcs;

llvm/test/CodeGen/AArch64/addsub-24bit-imm.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -run-pass=aarch64-mi-peephole-opt -o - -mtriple=aarch64-unknown-linux -verify-machineinstrs %s | FileCheck %s
 
 # Main intention is to verify machine instructions have valid register classes.
 # Use of UBFM[W|X]ri is used as an arbitrary instruction that requires GPR[32|64]RegClass.
 # If the ADD/SUB optimization generates invalid register classes, this test will fail.
 ---
 name: addi
 body: |
   bb.0.entry:
     liveins: $w0
     ; CHECK-LABEL: name: addi
-    ; CHECK: [[COPY:%[0-9]+]]:gpr32common = COPY $w0
+    ; CHECK: liveins: $w0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr32common = COPY $w0
     ; CHECK-NEXT: [[ADDWri:%[0-9]+]]:gpr32sp = ADDWri [[COPY]], 273, 12
     ; CHECK-NEXT: [[ADDWri1:%[0-9]+]]:gpr32common = ADDWri [[ADDWri]], 3549, 0
     ; CHECK-NEXT: [[UBFMWri:%[0-9]+]]:gpr32 = UBFMWri [[ADDWri1]], 28, 31
     ; CHECK-NEXT: $w0 = COPY [[UBFMWri]]
     ; CHECK-NEXT: RET_ReallyLR implicit $w0
     %0:gpr32 = COPY $w0
     %1:gpr32 = MOVi32imm 1121757
     %2:gpr32 = ADDWrr %0, %1
     %3:gpr32 = UBFMWri %2, 28, 31
     $w0 = COPY %3
     RET_ReallyLR implicit $w0
 ...
 ---
 name: addl
 body: |
   bb.0.entry:
     liveins: $x0
     ; CHECK-LABEL: name: addl
-    ; CHECK: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
     ; CHECK-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri [[COPY]], 273, 12
     ; CHECK-NEXT: [[ADDXri1:%[0-9]+]]:gpr64common = ADDXri [[ADDXri]], 3549, 0
     ; CHECK-NEXT: [[UBFMXri:%[0-9]+]]:gpr64 = UBFMXri [[ADDXri1]], 28, 31
     ; CHECK-NEXT: $x0 = COPY [[UBFMXri]]
     ; CHECK-NEXT: RET_ReallyLR implicit $x0
     %0:gpr64 = COPY $x0
     %1:gpr32 = MOVi32imm 1121757
     %2:gpr64 = SUBREG_TO_REG 0, %1, %subreg.sub_32
     %3:gpr64 = ADDXrr %0, killed %2
     %4:gpr64 = UBFMXri %3, 28, 31
     $x0 = COPY %4
     RET_ReallyLR implicit $x0
 ...
 ---
 name: addl_negate
 body: |
   bb.0.entry:
     liveins: $x0
     ; CHECK-LABEL: name: addl_negate
-    ; CHECK: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64common = COPY $x0
     ; CHECK-NEXT: [[SUBXri:%[0-9]+]]:gpr64sp = SUBXri [[COPY]], 273, 12
     ; CHECK-NEXT: [[SUBXri1:%[0-9]+]]:gpr64common = SUBXri [[SUBXri]], 3549, 0
     ; CHECK-NEXT: [[UBFMXri:%[0-9]+]]:gpr64 = UBFMXri [[SUBXri1]], 28, 31
     ; CHECK-NEXT: $x0 = COPY [[UBFMXri]]
     ; CHECK-NEXT: RET_ReallyLR implicit $x0
     %0:gpr64 = COPY $x0
     %1:gpr64 = MOVi64imm -1121757
     %2:gpr64 = ADDXrr %0, killed %1
     %3:gpr64 = UBFMXri %2, 28, 31
     $x0 = COPY %3
     RET_ReallyLR implicit $x0
+...
+---
+name: add_xzr
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: add_xzr
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64 = MOVi64imm -2105098
+    ; CHECK-NEXT: [[ADDXrr:%[0-9]+]]:gpr64common = ADDXrr $xzr, [[MOVi64imm]]
+    ; CHECK-NEXT: [[MADDXrrr:%[0-9]+]]:gpr64 = MADDXrrr [[COPY]], [[COPY]], [[ADDXrr]]
+    ; CHECK-NEXT: $x0 = COPY [[MADDXrrr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %2:gpr64 = MOVi64imm -2105098
+    %4:gpr64common = ADDXrr $xzr, %2
+    %3:gpr64 = MADDXrrr %0, %0, %4
+    $x0 = COPY %3
+    RET_ReallyLR implicit $x0
+...
+---
+name: sub_xzr
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: sub_xzr
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64 = MOVi64imm -2105098
+    ; CHECK-NEXT: [[SUBXrr:%[0-9]+]]:gpr64common = SUBXrr $xzr, [[MOVi64imm]]
+    ; CHECK-NEXT: [[MADDXrrr:%[0-9]+]]:gpr64 = MADDXrrr [[COPY]], [[COPY]], [[SUBXrr]]
+    ; CHECK-NEXT: $x0 = COPY [[MADDXrrr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %2:gpr64 = MOVi64imm -2105098
+    %4:gpr64common = SUBXrr $xzr, %2
+    %3:gpr64 = MADDXrrr %0, %0, %4
+    $x0 = COPY %3
+    RET_ReallyLR implicit $x0
+...
+---
+name: adds_xzr
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: adds_xzr
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64 = MOVi64imm -2105098
+    ; CHECK-NEXT: [[ADDSXrr:%[0-9]+]]:gpr64common = ADDSXrr $xzr, [[MOVi64imm]], implicit-def $nzcv
+    ; CHECK-NEXT: [[MADDXrrr:%[0-9]+]]:gpr64 = MADDXrrr [[COPY]], [[COPY]], [[ADDSXrr]]
+    ; CHECK-NEXT: $x0 = COPY [[MADDXrrr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %2:gpr64 = MOVi64imm -2105098
+    %4:gpr64common = ADDSXrr $xzr, %2, implicit-def $nzcv
+    %3:gpr64 = MADDXrrr %0, %0, %4
+    $x0 = COPY %3
+    RET_ReallyLR implicit $x0
+...
+---
+name: subs_xzr
+body: |
+  bb.0.entry:
+    liveins: $x0
+    ; CHECK-LABEL: name: subs_xzr
+    ; CHECK: liveins: $x0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64 = COPY $x0
+    ; CHECK-NEXT: [[MOVi64imm:%[0-9]+]]:gpr64 = MOVi64imm -2105098
+    ; CHECK-NEXT: [[SUBSXrr:%[0-9]+]]:gpr64common = SUBSXrr $xzr, [[MOVi64imm]], implicit-def $nzcv
+    ; CHECK-NEXT: [[MADDXrrr:%[0-9]+]]:gpr64 = MADDXrrr [[COPY]], [[COPY]], [[SUBSXrr]]
+    ; CHECK-NEXT: $x0 = COPY [[MADDXrrr]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64 = COPY $x0
+    %2:gpr64 = MOVi64imm -2105098
+    %4:gpr64common = SUBSXrr $xzr, %2, implicit-def $nzcv
+    %3:gpr64 = MADDXrrr %0, %0, %4
+    $x0 = COPY %3
+    RET_ReallyLR implicit $x0