[LLVM][AArch64ExpandPseudo] Preserve undef flags when expanding SVE 1/2/3-op pseudo instructions. (#149104)

Fixes #149034

Author: paulwalker-arm

llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp

@@ -598,6 +598,9 @@ bool AArch64ExpandPseudo::expand_DestructiveOp(
     llvm_unreachable("Unsupported ElementSize");
   }
 
+  // Preserve undef state until DOP's reg is defined.
+  unsigned DOPRegState = MI.getOperand(DOPIdx).isUndef() ? RegState::Undef : 0;
+
   //
   // Create the destructive operation (if required)
   //
@@ -616,10 +619,11 @@ bool AArch64ExpandPseudo::expand_DestructiveOp(
     PRFX = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(MovPrfxZero))
                .addReg(DstReg, RegState::Define)
                .addReg(MI.getOperand(PredIdx).getReg())
-               .addReg(MI.getOperand(DOPIdx).getReg());
+               .addReg(MI.getOperand(DOPIdx).getReg(), DOPRegState);
 
     // After the movprfx, the destructive operand is same as Dst
     DOPIdx = 0;
+    DOPRegState = 0;
 
     // Create the additional LSL to zero the lanes when the DstReg is not
     // unique. Zeros the lanes in z0 that aren't active in p0 with sequence
@@ -638,19 +642,21 @@ bool AArch64ExpandPseudo::expand_DestructiveOp(
     assert(DOPRegIsUnique && "The destructive operand should be unique");
     PRFX = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(MovPrfx))
                .addReg(DstReg, RegState::Define)
-               .addReg(MI.getOperand(DOPIdx).getReg());
+               .addReg(MI.getOperand(DOPIdx).getReg(), DOPRegState);
     DOPIdx = 0;
+    DOPRegState = 0;
   }
 
   //
   // Create the destructive operation
   //
   DOP = BuildMI(MBB, MBBI, MI.getDebugLoc(), TII->get(Opcode))
     .addReg(DstReg, RegState::Define | getDeadRegState(DstIsDead));
+  DOPRegState = DOPRegState | RegState::Kill;
 
   switch (DType) {
   case AArch64::DestructiveUnaryPassthru:
-    DOP.addReg(MI.getOperand(DOPIdx).getReg(), RegState::Kill)
+    DOP.addReg(MI.getOperand(DOPIdx).getReg(), DOPRegState)
         .add(MI.getOperand(PredIdx))
         .add(MI.getOperand(SrcIdx));
     break;
@@ -659,12 +665,12 @@ bool AArch64ExpandPseudo::expand_DestructiveOp(
   case AArch64::DestructiveBinaryComm:
   case AArch64::DestructiveBinaryCommWithRev:
     DOP.add(MI.getOperand(PredIdx))
-       .addReg(MI.getOperand(DOPIdx).getReg(), RegState::Kill)
-       .add(MI.getOperand(SrcIdx));
+        .addReg(MI.getOperand(DOPIdx).getReg(), DOPRegState)
+        .add(MI.getOperand(SrcIdx));
     break;
   case AArch64::DestructiveTernaryCommWithRev:
     DOP.add(MI.getOperand(PredIdx))
-        .addReg(MI.getOperand(DOPIdx).getReg(), RegState::Kill)
+        .addReg(MI.getOperand(DOPIdx).getReg(), DOPRegState)
         .add(MI.getOperand(SrcIdx))
         .add(MI.getOperand(Src2Idx));
     break;

llvm/test/CodeGen/AArch64/sve-pseudos-expand-undef.mir

@@ -54,3 +54,84 @@ body:             |
     renamable $z0 = FADD_ZPZZ_D_UNDEF killed $p0, killed $z1, killed $z2, implicit-def $z0_z1_z2_z3
     RET_ReallyLR implicit $z0_z1_z2_z3
 ...
+
+---
+name: unary_undef_operand
+body:             |
+  bb.0:
+    liveins: $p0, $z0
+
+    ; CHECK: name: unary_undef_operand
+    ; CHECK: $z0 = MOVPRFX_ZZ undef $z1
+    ; CHECK: $z0 = ABS_ZPmZ_S internal killed $z0, renamable $p0, killed undef renamable $z1
+    ; NOTE: Unary _UNDEF psuedo instructions ignore the passthru operand.
+    renamable $z0 = ABS_ZPmZ_S_UNDEF renamable $z0, renamable $p0, killed undef renamable $z1
+    RET_ReallyLR
+
+...
+
+---
+name: binop_undef_operand
+body:             |
+  bb.0:
+    liveins: $p0, $z1
+
+    ; CHECK: name: binop_undef_operand
+    ; CHECK-NOT: MOVPRFX
+    ; CHECK: $z0 = SMIN_ZPmZ_S renamable $p0, killed undef $z0, killed renamable $z1
+    renamable $z0 = SMIN_ZPZZ_S_UNDEF renamable $p0, undef renamable $z0, killed renamable $z1
+    RET_ReallyLR
+
+...
+
+---
+name: binop_undef_operand_requires_movpfrx
+body:             |
+  bb.0:
+    liveins: $p0, $z1
+
+    ; CHECK: name: binop_undef_operand_requires_movpfrx
+    ; CHECK: $z0 = MOVPRFX_ZZ undef $z2
+    ; CHECK: $z0 = SMIN_ZPmZ_S renamable $p0, internal killed $z0, killed renamable $z1
+    renamable $z0 = SMIN_ZPZZ_S_UNDEF renamable $p0, undef renamable $z2, killed renamable $z1
+    RET_ReallyLR
+
+...
+
+---
+name: binop_undef_operand_requires_zeroing_movpfrx
+body:             |
+  bb.0:
+    liveins: $p0, $z1
+
+    ; CHECK: name: binop_undef_operand_requires_zeroing_movpfrx
+    ; CHECK: $z0 = MOVPRFX_ZPzZ_S $p0, undef $z2
+    ; CHECK: $z0 = ADD_ZPmZ_S renamable $p0, internal killed $z0, killed renamable $z1
+    renamable $z0 = ADD_ZPZZ_S_ZERO renamable $p0, undef renamable $z2, killed renamable $z1
+    RET_ReallyLR
+
+...
+
+---
+name: ternaryop_undef_operand
+body:             |
+  bb.0:
+    liveins: $p0, $z1, $z2
+    ; CHECK: name: ternaryop_undef_operand
+    ; CHECK-NOT: MOVPRFX
+    ; CHECK: $z0 = MLA_ZPmZZ_B killed renamable $p0, killed undef $z0, killed renamable $z1, killed renamable $z2
+    renamable $z0 = MLA_ZPZZZ_B_UNDEF killed renamable $p0, killed undef renamable $z0, killed renamable $z1, killed renamable $z2
+    RET_ReallyLR implicit $z0
+...
+
+---
+name: ternaryop_undef_operand_requires_movprfx
+body:             |
+  bb.0:
+    liveins: $p0, $z1, $z2
+    ; CHECK: name: ternaryop_undef_operand_requires_movprfx
+    ; CHECK: $z0 = MOVPRFX_ZZ undef $z3
+    ; CHECK: $z0 = MLA_ZPmZZ_B killed renamable $p0, internal killed $z0, killed renamable $z1, killed renamable $z2
+    renamable $z0 = MLA_ZPZZZ_B_UNDEF killed renamable $p0, killed undef renamable $z3, killed renamable $z1, killed renamable $z2
+    RET_ReallyLR implicit $z0
+...