[AArch64][SVE] Fix missed immediate selection due to mishandling of signedness

[bsmith]

Summary

The complex selection pattern for add/sub shifted immediates is
incorrect in it's handling of incoming constant values, in that it
does not properly anticipate the values to be signed extended to
32-bits.

Co-authored-by: Graham Hunter <[e-mail]>

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3121

Does checking ((int32_t)ImmVal >= -32768) && ((int32_t)ImmVal <= 32512)) actually do anything? All 16-bit immediates should be in that range.

[bsmith]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3121

It possibly doesn't, however the thing that makes me nervous here is that the incoming CNode value is an i32 so could in theory (if something has gone wrong?) have a value outside of that range. It might be best to do it like this if nothing else to prevent an issue further up the pipeline propagating into a codegen fault here?

[david-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3121

I guess you could possibly remove those checks and replace them with an assert?

[bsmith]

• Fix clang format issues
• Move range check to be an assert since it should always be true

[peterwaller-arm]

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

3122

Suggestion: Use shift operation since this is about bit manipulation, and it's consistent with other uses nearby.

[bsmith]

• Use shift in place of divide for consistency.

llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp

if (((Imm % 256) == 0) && (Imm >= -32768) && (Imm <= 32512)) {
     return true;
   }
 
   return false;
 }
 
 bool AArch64DAGToDAGISel::SelectSVEAddSubImm(SDValue N, MVT VT, SDValue &Imm, SDValue &Shift) {
   if (auto CNode = dyn_cast<ConstantSDNode>(N)) {
-    const int64_t ImmVal = CNode->getZExtValue();
+    const int64_t ImmVal = CNode->getSExtValue();
     SDLoc DL(N);
 
     switch (VT.SimpleTy) {
     case MVT::i8:
+      // Can always select i8s, no shift, mask the immediate value to
+      // deal with sign-extended value from lowering.
+      Shift = CurDAG->getTargetConstant(0, DL, MVT::i32);
+      Imm = CurDAG->getTargetConstant(ImmVal & 0xFF, DL, MVT::i32);
+      return true;
+    case MVT::i16:
+      // i16 values get sign-extended to 32-bits during lowering.
       if ((ImmVal & 0xFF) == ImmVal) {
         Shift = CurDAG->getTargetConstant(0, DL, MVT::i32);
         Imm = CurDAG->getTargetConstant(ImmVal, DL, MVT::i32);
         return true;
+      } else if ((ImmVal & 0xFF) == 0) {
+        assert((ImmVal >= -32768) && (ImmVal <= 32512));
+        Shift = CurDAG->getTargetConstant(8, DL, MVT::i32);
+        Imm = CurDAG->getTargetConstant((ImmVal >> 8) & 0xFF, DL, MVT::i32);
+        return true;
       }
       break;
-    case MVT::i16:
     case MVT::i32:
     case MVT::i64:
+      // Range of immediate won't trigger signedness problems for 32/64b.
       if ((ImmVal & 0xFF) == ImmVal) {
         Shift = CurDAG->getTargetConstant(0, DL, MVT::i32);
         Imm = CurDAG->getTargetConstant(ImmVal, DL, MVT::i32);
         return true;
       } else if ((ImmVal & 0xFF00) == ImmVal) {
         Shift = CurDAG->getTargetConstant(8, DL, MVT::i32);
         Imm = CurDAG->getTargetConstant(ImmVal >> 8, DL, MVT::i32);
         return true;

llvm/test/CodeGen/AArch64/sve-int-imm.ll

 ; CHECK: add  z0.d, z0.d, #1024
 ; CHECK-NEXT: ret
   %elt = insertelement <vscale x 2 x i64> undef, i64 1024, i32 0
   %splat = shufflevector <vscale x 2 x i64> %elt, <vscale x 2 x i64> undef, <vscale x 2 x i32> zeroinitializer
   %res = add <vscale x 2 x i64> %a, %splat
   ret <vscale x 2 x i64> %res
 }
 
+define <vscale x 16 x i8> @add_i8_signedness(<vscale x 16 x i8> %a) {
+; CHECK-LABEL: add_i8_signedness
+; CHECK: add  z0.b, z0.b, #255
+; CHECK-NEXT: ret
+  %elt = insertelement <vscale x 16 x i8> undef, i8 255, i32 0
+  %splat = shufflevector <vscale x 16 x i8> %elt, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
+  %res =  add <vscale x 16 x i8> %a, %splat
+  ret <vscale x 16 x i8> %res
+}
+
+define <vscale x 8 x i16> @add_i16_signedness(<vscale x 8 x i16> %a) {
+; CHECK-LABEL: add_i16_signedness
+; CHECK: add  z0.h, z0.h, #65280
+; CHECK-NEXT: ret
+  %elt = insertelement <vscale x 8 x i16> undef, i16 65280, i32 0
+  %splat = shufflevector <vscale x 8 x i16> %elt, <vscale x 8 x i16> undef, <vscale x 8 x i32> zeroinitializer
+  %res =  add <vscale x 8 x i16> %a, %splat
+  ret <vscale x 8 x i16> %res
+}
+
 ; SUBR
 define <vscale x 16 x i8> @subr_i8_low(<vscale x 16 x i8> %a) {
 ; CHECK-LABEL: subr_i8_low
 ; CHECK: subr  z0.b, z0.b, #30
 ; CHECK-NEXT: ret
   %elt = insertelement <vscale x 16 x i8> undef, i8 30, i32 0
   %splat = shufflevector <vscale x 16 x i8> %elt, <vscale x 16 x i8> undef, <vscale x 16 x i32> zeroinitializer
   %res =  sub <vscale x 16 x i8> %splat, %a