[AArch64] Emit zext move when the source of the zext is AssertZext or AssertSext

[wwei]

Summary

When the source of the zext is AssertZext or AssertSext, it is hard to know any information about the upper 32 bits,
so we should insert a zext move before emitting SUBREG_TO_REG to define the lower 32 bits.

Bugzilla: https://bugs.llvm.org/show_bug.cgi?id=47543

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

428

AssertSext/AssertZext don't really indicate anything either way; it would make sense to check the operand, maybe? That said, I guess the operand will usually be a CopyFromReg, so maybe it doesn't matter much.

More generally, I don't like the general approach of guessing what isel will do, as opposed to examining what isel actually did. Due to the way the dot product patterns were written, it's actually possible for an i32 EXTRACT_VECTOR_ELT to produce a value that isn't zero-extended (testcase follows). And I'm not confident there aren't other weird edge cases. I'd be happier doing this after isel, when we can tell what instruction actually produced the value in question.

define i64 @test_udot_v8i8(i8* nocapture readonly %a, i8* nocapture readonly %b) {
entry:
; CHECK-LABEL: test_udot_v8i8:
; CHECK:  udot {{v[0-9]+}}.2s, {{v[0-9]+}}.8b, {{v[0-9]+}}.8b
  %0 = bitcast i8* %a to <8 x i8>*
  %1 = load <8 x i8>, <8 x i8>* %0
  %2 = zext <8 x i8> %1 to <8 x i32>
  %3 = bitcast i8* %b to <8 x i8>*
  %4 = load <8 x i8>, <8 x i8>* %3
  %5 = zext <8 x i8> %4 to <8 x i32>
  %6 = mul nuw nsw <8 x i32> %5, %2
  %7 = call i32 @llvm.experimental.vector.reduce.add.v8i32(<8 x i32> %6)
  %8 = zext i32 %7 to i64
  ret i64 %8
}
declare i32 @llvm.experimental.vector.reduce.add.v8i32(<8 x i32>)

llvm/test/CodeGen/AArch64/arm64-assert-zext-sext.ll

1 ↗

(On Diff #292236)

This testcase is way too complicated; can you extract out the bit that actually triggers this issue?

llvm/test/CodeGen/AArch64/shift_minsize.ll

62

This mov shouldn't be necessary, but the reason isn't really related to this patch, I guess.

[wwei]

llvm/test/CodeGen/AArch64/arm64-assert-zext-sext.ll

1 ↗

(On Diff #292236)

Done， test cases have been simplified，

[wwei]

llvm/lib/Target/AArch64/AArch64ISelLowering.h

428

I basically agree with you, it seems more reasonable doing it after isel.
Especially when DAG isel is based on basic block scope, it is difficult to consider global data/control flow. For example，in arm64-assert-zext-sext.ll case, both AssertZext and AssertZext will have operand CopyFromReg, but what really matters is where the operand (maybe EXTRACT_SUBREG or Truncate) of CopyFromReg comes from. Unfortunately, they're from other blocks. Therefore, in the ISEL process, we can only perform conservative instruction selection. This is the right way to fix bugs in pr47543.
I have also considered adding check for the operand of AssertSext/AssertZext, like below:

static inline bool isDef32(const SDNode &N) {
  unsigned Opc = N.getOpcode();
  if(Opc == AssertSext || Opc == AssertZext)
    Opc = N.getOperand(0).getOpcode();
  return Opc != ISD::TRUNCATE && Opc != TargetOpcode::EXTRACT_SUBREG &&
         Opc != ISD::CopyFromReg;
}

But I'm not sure whether the check is required. Finally, I refer to the implementation in x86.

def def32 : PatLeaf<(i32 GR32:$src), [{
  return N->getOpcode() != ISD::TRUNCATE &&
         N->getOpcode() != TargetOpcode::EXTRACT_SUBREG &&
         N->getOpcode() != ISD::CopyFromReg &&
         N->getOpcode() != ISD::AssertSext &&
         N->getOpcode() != ISD::AssertZext;
}]>;

Doing it after isel will be better. However, many scenarios may need to be considered. It is not clear what edge or strange cases may occur. The current patch is used as a trade-off solution.

[efriedma]

LGTM

I'm okay with taking this as an incremental improvement, and then looking into a different approach later.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

428

I don't anticipate any complex issues performing the transform after isel, but I could be missing something.

llvm/lib/Target/AArch64/AArch64ISelLowering.h

 
 } // end namespace AArch64ISD
 
 namespace {
 
 // Any instruction that defines a 32-bit result zeros out the high half of the
 // register. Truncate can be lowered to EXTRACT_SUBREG. CopyFromReg may
 // be copying from a truncate. But any other 32-bit operation will zero-extend
-// up to 64 bits.
+// up to 64 bits. AssertSext/AssertZext aren't saying anything about the upper
+// 32 bits, they're probably just qualifying a CopyFromReg.
 // FIXME: X86 also checks for CMOV here. Do we need something similar?
 static inline bool isDef32(const SDNode &N) {
   unsigned Opc = N.getOpcode();
   return Opc != ISD::TRUNCATE && Opc != TargetOpcode::EXTRACT_SUBREG &&
-         Opc != ISD::CopyFromReg;
+         Opc != ISD::CopyFromReg && Opc != ISD::AssertSext &&
+         Opc != ISD::AssertZext;
 }
 
 } // end anonymous namespace
 
 class AArch64Subtarget;
 class AArch64TargetMachine;
 
 class AArch64TargetLowering : public TargetLowering {

llvm/test/CodeGen/AArch64/shift_minsize.ll

 }
 
 define dso_local { i64, i64 } @shl128(i64 %x.coerce0, i64 %x.coerce1, i8 signext %y) minsize optsize {
 ; CHECK-LABEL: shl128:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    // kill: def $w2 killed $w2 def $x2
+; CHECK-NEXT:    mov w2, w2
 ; CHECK-NEXT:    bl __ashlti3
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 
 entry:
   %x.sroa.2.0.insert.ext = zext i64 %x.coerce1 to i128
   %x.sroa.2.0.insert.shift = shl nuw i128 %x.sroa.2.0.insert.ext, 64
   %x.sroa.0.0.insert.ext = zext i64 %x.coerce0 to i128

 }
 
 define dso_local { i64, i64 } @ashr128(i64 %x.coerce0, i64 %x.coerce1, i8 signext %y) minsize optsize {
 ; CHECK-LABEL: ashr128:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    // kill: def $w2 killed $w2 def $x2
+; CHECK-NEXT:    mov w2, w2
 ; CHECK-NEXT:    bl __ashrti3
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 entry:
   %x.sroa.2.0.insert.ext = zext i64 %x.coerce1 to i128
   %x.sroa.2.0.insert.shift = shl nuw i128 %x.sroa.2.0.insert.ext, 64
   %x.sroa.0.0.insert.ext = zext i64 %x.coerce0 to i128
   %x.sroa.0.0.insert.insert = or i128 %x.sroa.2.0.insert.shift, %x.sroa.0.0.insert.ext

 }
 
 define dso_local { i64, i64 } @lshr128(i64 %x.coerce0, i64 %x.coerce1, i8 signext %y) minsize optsize {
 ; CHECK-LABEL: lshr128:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    .cfi_offset w30, -16
-; CHECK-NEXT:    // kill: def $w2 killed $w2 def $x2
+; CHECK-NEXT:    mov w2, w2
 ; CHECK-NEXT:    bl __lshrti3
 ; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
 ; CHECK-NEXT:    ret
 entry:
   %x.sroa.2.0.insert.ext = zext i64 %x.coerce1 to i128
   %x.sroa.2.0.insert.shift = shl nuw i128 %x.sroa.2.0.insert.ext, 64
   %x.sroa.0.0.insert.ext = zext i64 %x.coerce0 to i128
   %x.sroa.0.0.insert.insert = or i128 %x.sroa.2.0.insert.shift, %x.sroa.0.0.insert.ext