[SelectionDAG] PromoteIntRes_EXTRACT_SUBVECTOR for scalable vectors (widening).

[sdesmalen]

Summary

This patch implements legalization of EXTRACT_SUBVECTOR for the case
where the result needs promoting, and the input type requires widening.

[frasercrmck]

Do you think it'd be possible to add RVV tests too? vscale x 2 doesn't need widening but vscale x 3 or something would.

[sdesmalen]

Duplicated tests.

[sdesmalen]

In D109509#2991655, @frasercrmck wrote:

Do you think it'd be possible to add RVV tests too? vscale x 2 doesn't need widening but vscale x 3 or something would.

Sure np. I've copied the same tests that I added for AArch64 and re-ran the update-test-checks script.

[frasercrmck]

In D109509#2994149, @sdesmalen wrote:

In D109509#2991655, @frasercrmck wrote:

Do you think it'd be possible to add RVV tests too? vscale x 2 doesn't need widening but vscale x 3 or something would.

Sure np. I've copied the same tests that I added for AArch64 and re-ran the update-test-checks script.

Thanks! But <vscale x 2 x i8> isn't illegal for us so, if I'm not mistaken, it's not hitting this particular promotion path.

[sdesmalen]

In D109509#2994154, @frasercrmck wrote:

In D109509#2994149, @sdesmalen wrote:

In D109509#2991655, @frasercrmck wrote:

Do you think it'd be possible to add RVV tests too? vscale x 2 doesn't need widening but vscale x 3 or something would.

Sure np. I've copied the same tests that I added for AArch64 and re-ran the update-test-checks script.

Thanks! But <vscale x 2 x i8> isn't illegal for us so, if I'm not mistaken, it's not hitting this particular promotion path.

Ah, sorry I misunderstood that!

What result type would require promotion? I'd need an example where the result would need to be promoted and the input widened.

[frasercrmck]

In D109509#2994157, @sdesmalen wrote:

Ah, sorry I misunderstood that!

What result type would require promotion? I'd need an example where the result would need to be promoted and the input widened.

My bad, I should have been clearer. I hope I wasn't too hasty in my suggestion; I was initially thinking of something like <vscale x 3 x i4> from <vscale x 9 x i4> but they both need promoting and widening, and now I've forgotten in which order those take place.

[sdesmalen]

In D109509#2994170, @frasercrmck wrote:

In D109509#2994157, @sdesmalen wrote:

Ah, sorry I misunderstood that!

What result type would require promotion? I'd need an example where the result would need to be promoted and the input widened.

My bad, I should have been clearer. I hope I wasn't too hasty in my suggestion; I was initially thinking of something like <vscale x 3 x i4> from <vscale x 9 x i4> but they both need promoting and widening, and now I've forgotten in which order those take place.

It seems that trying with an example of i4's (e.g. extracting nxv2i4 from nxv12i4), it actually runs into various other failures that would need fixing first, e.g. widening the truncate to i4's of the input vector. The exact example you gave above (where both need widening), actually fails even earlier in SelectionDAGBuilder::getCopyToPartsVector.

I'm not sure I can come up with an example where it does go down this code path though. Do you want me to keep the tests even if they don't exactly test the code in this patch?

[frasercrmck]

In D109509#2994216, @sdesmalen wrote:

It seems that trying with an example of i4's (e.g. extracting nxv2i4 from nxv12i4), it actually runs into various other failures that would need fixing first, e.g. widening the truncate to i4's of the input vector. The exact example you gave above (where both need widening), actually fails even earlier in SelectionDAGBuilder::getCopyToPartsVector.

I'm not sure I can come up with an example where it does go down this code path though. Do you want me to keep the tests even if they don't exactly test the code in this patch?

Ah yes I see, in the former the truncate is generated during argument lowering. That's a shame.

Then no, I don't think I need these tests because widening the EXTRACT_SUBVECTOR operand is trivial.

Sorry for wasting your time, @sdesmalen. I think this patch LGTM.

[sdesmalen]

In D109509#2994261, @frasercrmck wrote:

In D109509#2994216, @sdesmalen wrote:

It seems that trying with an example of i4's (e.g. extracting nxv2i4 from nxv12i4), it actually runs into various other failures that would need fixing first, e.g. widening the truncate to i4's of the input vector. The exact example you gave above (where both need widening), actually fails even earlier in SelectionDAGBuilder::getCopyToPartsVector.

I'm not sure I can come up with an example where it does go down this code path though. Do you want me to keep the tests even if they don't exactly test the code in this patch?

Ah yes I see, in the former the truncate is generated during argument lowering. That's a shame.

Then no, I don't think I need these tests because widening the EXTRACT_SUBVECTOR operand is trivial.

Sorry for wasting your time, @sdesmalen. I think this patch LGTM.

No worries, thanks for having a look! I'll remove the tests before committing.

llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp

if (getTypeAction(InVT) == TargetLowering::TypeSplitVector ||
                                   DAG.getConstant(alignDown(IdxVal, NElts), dl,
                                                   BaseIdx.getValueType()));
       SDValue Step2 = DAG.getNode(
           ISD::EXTRACT_SUBVECTOR, dl, OutVT, Step1,
           DAG.getConstant(IdxVal % NElts, dl, BaseIdx.getValueType()));
       return DAG.getNode(ISD::ANY_EXTEND, dl, NOutVT, Step2);
     }
 
+    // Try and extract from a widened type.
+    if (getTypeAction(InVT) == TargetLowering::TypeWidenVector) {
+      SDValue Ops[] = {GetWidenedVector(InOp0), BaseIdx};
+      SDValue Ext = DAG.getNode(ISD::EXTRACT_SUBVECTOR, SDLoc(N), OutVT, Ops);
+      return DAG.getNode(ISD::ANY_EXTEND, dl, NOutVT, Ext);
+    }
+
     // Promote operands and see if this is handled by target lowering,
     // Otherwise, use the BUILD_VECTOR approach below
     if (getTypeAction(InVT) == TargetLowering::TypePromoteInteger) {
       // Collect the (promoted) operands
       SDValue Ops[] = { GetPromotedInteger(InOp0), BaseIdx };
 
       EVT PromEltVT = Ops[0].getValueType().getVectorElementType();
       assert(PromEltVT.bitsLE(NOutVTElem) &&

llvm/test/CodeGen/AArch64/sve-extract-vector.ll

 ; CHECK-NEXT:    uunpkhi z0.h, z1.b
 ; CHECK-NEXT:    uunpkhi z0.s, z0.h
 ; CHECK-NEXT:    uunpkhi z0.d, z0.s
 ; CHECK-NEXT:    ret
   %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv32i8(<vscale x 32 x i8> %vec, i64 30)
   ret <vscale x 2 x i8> %res
 }
 
+;
+; Extracting illegal vector that needs promotion from a vector that needs widening.
+;
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_0(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_0:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.h, z0.b
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    uunpklo z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 0)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_2(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.h, z0.b
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 2)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_4(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_4:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.h, z0.b
+; CHECK-NEXT:    uunpkhi z0.s, z0.h
+; CHECK-NEXT:    uunpklo z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 4)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_6(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_6:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpklo z0.h, z0.b
+; CHECK-NEXT:    uunpkhi z0.s, z0.h
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 6)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_8(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpkhi z0.h, z0.b
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    uunpklo z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 8)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_10(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_10:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpkhi z0.h, z0.b
+; CHECK-NEXT:    uunpklo z0.s, z0.h
+; CHECK-NEXT:    uunpkhi z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 10)
+  ret <vscale x 2 x i8> %res
+}
+
+define <vscale x 2 x i8> @extract_nxv2i8_nxv14i8_12(<vscale x 14 x i8> %vec) {
+; CHECK-LABEL: extract_nxv2i8_nxv14i8_12:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uunpkhi z0.h, z0.b
+; CHECK-NEXT:    uunpkhi z0.s, z0.h
+; CHECK-NEXT:    uunpklo z0.d, z0.s
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> %vec, i64 12)
+  ret <vscale x 2 x i8> %res
+}
+
 attributes #0 = { vscale_range(2,2) }
 
 declare <2 x i64> @llvm.experimental.vector.extract.v2i64.nxv2i64(<vscale x 2 x i64>, i64)
 declare <4 x i32> @llvm.experimental.vector.extract.v4i32.nxv4i32(<vscale x 4 x i32>, i64)
 declare <8 x i16> @llvm.experimental.vector.extract.v8i16.nxv8i16(<vscale x 8 x i16>, i64)
 declare <16 x i8> @llvm.experimental.vector.extract.v16i8.nxv16i8(<vscale x 16 x i8>, i64)
 
 declare <4 x i64> @llvm.experimental.vector.extract.v4i64.nxv2i64(<vscale x 2 x i64>, i64)
 
 declare <vscale x 1 x i32> @llvm.experimental.vector.extract.nxv1i32.nxv4i32(<vscale x 4 x i32>, i64)
 declare <vscale x 1 x i16> @llvm.experimental.vector.extract.nxv1i16.nxv6i16(<vscale x 6 x i16>, i64)
 
 declare <vscale x 2 x i1> @llvm.experimental.vector.extract.nxv2i1.nxv16i1(<vscale x 16 x i1>, i64)
 declare <vscale x 2 x i1> @llvm.experimental.vector.extract.nxv2i1.nxv8i1(<vscale x 8 x i1>, i64)
 declare <vscale x 4 x i1> @llvm.experimental.vector.extract.nxv4i1.nxv16i1(<vscale x 16 x i1>, i64)
 
 declare <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv32i8(<vscale x 32 x i8> , i64)
+declare <vscale x 2 x i8> @llvm.experimental.vector.extract.nxv2i8.nxv14i8(<vscale x 14 x i8> , i64)