[LegalizeTypes][AArch64] Use scalar_to_vector to eliminate bitcast

[Allen]

Summary

Legalize t3: v2i16 = bitcast i32
with            (v2i16 extract_subvector (v4i16 bitcast (v2i32 scalar_to_vector (i32 in))), 0)

Fix https://github.com/llvm/llvm-project/issues/61638

NOTE: Don't touch getPreferredVectorAction like X86 as this will touch too many test cases.

[paulwalker-arm]

My gut feeling is this isn't a valid transformation. Looking at the output from:

define <2 x i16> @bitcast_from_int(i32 %word) {
  %ret = bitcast i32 %word to <2 x i16>
  ret <2 x i16> %ret
}

define i16 @extract_lo(<2 x i16> %vec) {
  %ret = extractelement <2 x i16> %vec, i64 0
  ret i16 %ret
}

define i16 @extract_hi(<2 x i16> %vec) {
  %ret = extractelement <2 x i16> %vec, i64 1
  ret i16 %ret
}

bitcast_from_int:                       // @bitcast_from_int
	sub	sp, sp, #16
	.cfi_def_cfa_offset 16
	add	x8, sp, #12
	str	w0, [sp, #12]
	ld1	{ v0.h }[0], [x8]
	orr	x8, x8, #0x2
	ld1	{ v0.h }[2], [x8]
	add	sp, sp, #16
	ret

extract_lo:                             // @extract_lo
	fmov	w0, s0
	ret

extract_hi:                             // @extract_hi
	mov	w0, v0.s[1]
	ret

Shows <2 x i16> is represented as an unpacked <2 x i32> vector with the valid bits being the bottom 16-bits within each 32bit element. This means the ISD::ANY_EXTEND within the DAG your matching is critical and this gets dropped when replacing with ISD::SCALAR_TO_VECTOR. Ultimately it looks like ISD::BITCAST for smaller than legal vectors types is not a nop unless the src/dst is a load/store in which case the extension can be removed. The poor code is the result of type legalisation so perhaps you'd be better off implementing custom legalisation for ISD::BITCAST involving 32-bit vector types. I recall we did something similar for SVE to handle its unpacked vector types (see ReplaceBITCASTResults).

[Allen]

Thanks for your advise, I'll try that.

• In current case, I also check N->getValueType(0).getScalarType() == N00.getNode()->getValueType(0), so we only address the case A --> VecB --> VecA, as the scalar type of VecA is same to A, so we can ignore the difference between VecB and VecA on the laid out ?

• In addition, do you mean we “ prefer to widen v4i16 to v2i32” similar to x86 as following ? (I test it , it works fine, but there is too many case need to update, so I want to make sure I have understood what you mean rightly for off implementing custom legalisation for ISD::BITCAST)

@@ -22680,6 +22680,11 @@ AArch64TargetLowering::getPreferredVectorAction(MVT VT) const {
       VT == MVT::v1f32)
     return TypeWidenVector;
 
+  // We prefer to widen v4i16 to v2i32 instead of to promote.
+  // TODO: Add more types.
+  if (VT == MVT::v2i16)
+    return TypeWidenVector;
+
   return TargetLoweringBase::getPreferredVectorAction(VT);
 }

[Allen]

Define custom legalisation for ISD::BITCAST as comment

[efriedma]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

22895

SCALAR_TO_VECTOR from i32 to v2i16 implicitly discards the high 16 bits of the input. (The other lane is UNDEF. UNDEF might happen to be what you want in simple cases, but it won't be in general.)

Something like (v2i16 extract_subvector (v4i16 bitcast (v2i32 scalar_to_vector (i32 in))), 0) should have the semantics you want, and lower to something reasonably efficient.

[Allen]

(v2i16 extract_subvector (v4i16 bitcast (v2i32 scalar_to_vector (i32 in))), 0)

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

22895

Thanks, apply your comment.
The extra instruction ushll v0.4s, v0.4h, #0 for (v4i32 any_extend v4i16) will clear the other lane.

        t2: i32,ch = CopyFromReg t0, Register:i32 %0
      t8: v2i32 = scalar_to_vector t2
    t9: v4i16 = bitcast t8
  t20: v4i32 = any_extend t9
t21: v2i32 = extract_subvector t20, Constant:i64<0>

[dmgreen]

Is it worth adding other types like i16->v2i8 and i32->v4i8 too?

[Allen]

Add new biscast i16->v2i8 and i32->v4i8

[dmgreen]

Thanks for doing this. That's great.

The v2i8 version doesn't look too efficient with it moving back and forth, but they all look better than going via the stack. LGTM, thanks.

[Allen]

In D147678#4416884, @dmgreen wrote:

Thanks for doing this. That's great.

The v2i8 version doesn't look too efficient with it moving back and forth, but they all look better than going via the stack. LGTM, thanks.

I add a comment on https://github.com/llvm/llvm-project/issues/61638#issuecomment-1495316614** , these issue can produce efficient assemble (only a fmov s0, w0** as expected), but there is too many tests affected.
so I'm not sure if this is a direction for further optimization?

[paulwalker-arm]

I've an extra nit but otherwise looks good.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

22894–22897

Is this necessary given it is a requirement of ISD::BITCAST? By which I mean you shouldn't need to manually verify the DAG is valid here? I'd expect getNode() to catch such errors.

22902

Please can you use getVectorIdxConstant here rather than hardwiring MVT::i64?

[Allen]

Fix comments
a) delete assert
b) use getVectorIdxConstant

[Allen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

22894–22897

deleted assert, thanks

22902

Apply your comment, thanks

[dmgreen]

In D147678#4416899, @Allen wrote:

I add a comment on https://github.com/llvm/llvm-project/issues/61638#issuecomment-1495316614** , these issue can produce efficient assemble (only a fmov s0, w0** as expected), but there is too many tests affected.
so I'm not sure if this is a direction for further optimization?

The <2 x i16> across a call boundary needs to be treated like a <2 x i32> (as that is a legal type). To do differently would be an ABI break, which would be a lot of work. There are certainly places where we could be better for small vector types, but I don't think a fmov s0, w0 would be valid on it's own. My comment was more about the bitcast i16 %word to <2 x i8> case which is moving back and forth between gpr and vectors more than it needs to. That can be improved in other patches though.

By the way it might be worth making sure the tests cover BE too. I believe they are OK, but it would be a good idea to make sure it is tested.

[Allen]

rebase to cover BE

[Allen]

rebase as the changes of tests

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasNEON()) {
 
     setTruncStoreAction(MVT::v4i16, MVT::v4i8, Custom);
 
     setOperationAction(ISD::BITCAST, MVT::i2, Custom);
     setOperationAction(ISD::BITCAST, MVT::i4, Custom);
     setOperationAction(ISD::BITCAST, MVT::i8, Custom);
     setOperationAction(ISD::BITCAST, MVT::i16, Custom);
 
+    setOperationAction(ISD::BITCAST, MVT::v2i8, Custom);
+    setOperationAction(ISD::BITCAST, MVT::v2i16, Custom);
+    setOperationAction(ISD::BITCAST, MVT::v4i8, Custom);
+
     setLoadExtAction(ISD::EXTLOAD,  MVT::v4i16, MVT::v4i8, Custom);
     setLoadExtAction(ISD::SEXTLOAD, MVT::v4i16, MVT::v4i8, Custom);
     setLoadExtAction(ISD::ZEXTLOAD, MVT::v4i16, MVT::v4i8, Custom);
     setLoadExtAction(ISD::EXTLOAD,  MVT::v4i32, MVT::v4i8, Custom);
     setLoadExtAction(ISD::SEXTLOAD, MVT::v4i32, MVT::v4i8, Custom);
     setLoadExtAction(ISD::ZEXTLOAD, MVT::v4i32, MVT::v4i8, Custom);
 
     // ADDP custom lowering
for (SDNode *Node : Copy->uses()) {
     HasRet = true;
   }
 
   if (!HasRet)
     return false;
 
   Chain = TCChain;
   return true;
 }
 
 // Return whether the an instruction can potentially be optimized to a tail
 // call. This will cause the optimizers to attempt to move, or duplicate,
 // return instructions to help enable tail call optimizations for this
 // instruction.
 bool AArch64TargetLowering::mayBeEmittedAsTailCall(const CallInst *CI) const {
   return CI->isTailCall();
 }
 
 bool AArch64TargetLowering::getIndexedAddressParts(SDNode *N, SDNode *Op,
                                                    SDValue &Base,
                                                    SDValue &Offset,
                                                    SelectionDAG &DAG) const {
   if (Op->getOpcode() != ISD::ADD && Op->getOpcode() != ISD::SUB)
     return false;
 
if (AllUndef)
       Op = Op.getOperand(0);
   }
 
   SDValue VectorBits = vectorToScalarBitmask(Op.getNode(), DAG);
   if (VectorBits)
     Results.push_back(DAG.getZExtOrTrunc(VectorBits, DL, VT));
 }
 
+static void CustomNonLegalBITCASTResults(SDNode *N,
+                                         SmallVectorImpl<SDValue> &Results,
+                                         SelectionDAG &DAG, EVT ExtendVT,
+                                         EVT CastVT) {
+  SDLoc DL(N);
+  SDValue Op = N->getOperand(0);
+  EVT VT = N->getValueType(0);
+
+  // Use SCALAR_TO_VECTOR for lane zero
+  SDValue Vec = DAG.getNode(ISD::SCALAR_TO_VECTOR, DL, ExtendVT, Op);
+  SDValue CastVal = DAG.getNode(ISD::BITCAST, DL, CastVT, Vec);
+  SDValue IdxZero = DAG.getVectorIdxConstant(0, DL);
+  Results.push_back(
+      DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, VT, CastVal, IdxZero));
+  return;
+}
+
 void AArch64TargetLowering::ReplaceBITCASTResults(
     SDNode *N, SmallVectorImpl<SDValue> &Results, SelectionDAG &DAG) const {
   SDLoc DL(N);
   SDValue Op = N->getOperand(0);
   EVT VT = N->getValueType(0);
   EVT SrcVT = Op.getValueType();
 
+  if (VT == MVT::v2i16 && SrcVT == MVT::i32) {
+    CustomNonLegalBITCASTResults(N, Results, DAG, MVT::v2i32, MVT::v4i16);
+    return;
+  }
+
+  if (VT == MVT::v4i8 && SrcVT == MVT::i32) {
+    CustomNonLegalBITCASTResults(N, Results, DAG, MVT::v2i32, MVT::v8i8);
+    return;
+  }
+
+  if (VT == MVT::v2i8 && SrcVT == MVT::i16) {
+    CustomNonLegalBITCASTResults(N, Results, DAG, MVT::v4i16, MVT::v8i8);
+    return;
+  }
+
   if (VT.isScalableVector() && !isTypeLegal(VT) && isTypeLegal(SrcVT)) {
     assert(!VT.isFloatingPoint() && SrcVT.isFloatingPoint() &&
            "Expected fp->int bitcast!");
 
     // Bitcasting between unpacked vector types of different element counts is
     // not a NOP because the live elements are laid out differently.
     //                01234567
     // e.g. nxv2i32 = XX??XX??

llvm/test/CodeGen/AArch64/aarch64-load-ext.ll

; CHECK-BE-NEXT:    ret
   %z = add <4 x i8> %x, %y
   %s = sext <4 x i8> %z to <4 x i32>
   ret <4 x i32> %s
 }
 
 define <4 x i8> @bitcast(i32 %0) {
 ; CHECK-LE-LABEL: bitcast:
 ; CHECK-LE:       // %bb.0:
-; CHECK-LE-NEXT:    sub sp, sp, #16
-; CHECK-LE-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-LE-NEXT:    str w0, [sp, #12]
-; CHECK-LE-NEXT:    ldr s0, [sp, #12]
-; CHECK-LE-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-LE-NEXT:    // kill: def $d0 killed $d0 killed $q0
-; CHECK-LE-NEXT:    add sp, sp, #16
+; CHECK-LE-NEXT:    fmov s0, w0
+; CHECK-LE-NEXT:    zip1 v0.8b, v0.8b, v0.8b
 ; CHECK-LE-NEXT:    ret
 ;
 ; CHECK-BE-LABEL: bitcast:
 ; CHECK-BE:       // %bb.0:
-; CHECK-BE-NEXT:    sub sp, sp, #16
-; CHECK-BE-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-BE-NEXT:    str w0, [sp, #12]
-; CHECK-BE-NEXT:    ldr s0, [sp, #12]
+; CHECK-BE-NEXT:    fmov s0, w0
 ; CHECK-BE-NEXT:    rev32 v0.8b, v0.8b
-; CHECK-BE-NEXT:    ushll v0.8h, v0.8b, #0
-; CHECK-BE-NEXT:    rev64 v0.4h, v0.4h
-; CHECK-BE-NEXT:    add sp, sp, #16
+; CHECK-BE-NEXT:    zip1 v0.8b, v0.8b, v0.8b
+; CHECK-BE-NEXT:    rev64 v0.8b, v0.8b
 ; CHECK-BE-NEXT:    ret
   %2 = bitcast i32 %0 to <4 x i8>
   ret <4 x i8> %2
 }

llvm/test/CodeGen/AArch64/neon-bitcast.ll

 define <16 x i8> @test_v2f64_to_v16i8(<2 x double> %in) nounwind{
 ; CHECK-LABEL: test_v2f64_to_v16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ret
   %val = bitcast <2 x double> %in to <16 x i8>
   ret <16 x i8> %val
 }
 
+define <2 x i16> @bitcast_i32_to_v2i16(i32 %word) {
+; CHECK-LE-LABEL: bitcast_i32_to_v2i16:
+; CHECK-LE:       // %bb.0:
+; CHECK-LE-NEXT:    fmov s0, w0
+; CHECK-LE-NEXT:    ushll v0.4s, v0.4h, #0
+; CHECK-LE-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-LE-NEXT:    ret
+;
+; CHECK-BE-LABEL: bitcast_i32_to_v2i16:
+; CHECK-BE:       // %bb.0:
+; CHECK-BE-NEXT:    fmov s0, w0
+; CHECK-BE-NEXT:    rev32 v0.4h, v0.4h
+; CHECK-BE-NEXT:    ushll v0.4s, v0.4h, #0
+; CHECK-BE-NEXT:    rev64 v0.2s, v0.2s
+; CHECK-BE-NEXT:    ret
+  %ret = bitcast i32 %word to <2 x i16>
+  ret <2 x i16> %ret
+}
+
+define <4 x i8> @bitcast_i32_to_v4i8(i32 %word) {
+; CHECK-LE-LABEL: bitcast_i32_to_v4i8:
+; CHECK-LE:       // %bb.0:
+; CHECK-LE-NEXT:    fmov s0, w0
+; CHECK-LE-NEXT:    zip1 v0.8b, v0.8b, v0.8b
+; CHECK-LE-NEXT:    ret
+;
+; CHECK-BE-LABEL: bitcast_i32_to_v4i8:
+; CHECK-BE:       // %bb.0:
+; CHECK-BE-NEXT:    fmov s0, w0
+; CHECK-BE-NEXT:    rev32 v0.8b, v0.8b
+; CHECK-BE-NEXT:    zip1 v0.8b, v0.8b, v0.8b
+; CHECK-BE-NEXT:    rev64 v0.8b, v0.8b
+; CHECK-BE-NEXT:    ret
+  %ret = bitcast i32 %word to <4 x i8>
+  ret <4 x i8> %ret
+}
+
+; TODO: Eliminate redundant moving back and forth between gpr and vectors
+define <2 x i8> @bitcast_i16_to_v2i8(i16 %word) {
+; CHECK-LE-LABEL: bitcast_i16_to_v2i8:
+; CHECK-LE:       // %bb.0:
+; CHECK-LE-NEXT:    fmov s0, w0
+; CHECK-LE-NEXT:    umov w8, v0.b[0]
+; CHECK-LE-NEXT:    umov w9, v0.b[1]
+; CHECK-LE-NEXT:    fmov s0, w8
+; CHECK-LE-NEXT:    mov v0.s[1], w9
+; CHECK-LE-NEXT:    // kill: def $d0 killed $d0 killed $q0
+; CHECK-LE-NEXT:    ret
+;
+; CHECK-BE-LABEL: bitcast_i16_to_v2i8:
+; CHECK-BE:       // %bb.0:
+; CHECK-BE-NEXT:    fmov s0, w0
+; CHECK-BE-NEXT:    rev16 v0.16b, v0.16b
+; CHECK-BE-NEXT:    umov w8, v0.b[0]
+; CHECK-BE-NEXT:    umov w9, v0.b[1]
+; CHECK-BE-NEXT:    fmov s0, w8
+; CHECK-BE-NEXT:    mov v0.s[1], w9
+; CHECK-BE-NEXT:    rev64 v0.2s, v0.2s
+; CHECK-BE-NEXT:    ret
+  %ret = bitcast i16 %word to <2 x i8>
+  ret <2 x i8> %ret
+}
+