[AArch64][SelectionDAG] Refactor to support more scalable vector extending loads

[Allen]

Summary

Accord the discussion in D120953, we should firstly exclude all scalable vector
extending loads and then selectively enable those which we directly support.

This patch is intend to refactor for above (truncating stores is not touched),and
more scalable vector types will try to reduce the number of masked loads in favour
of more unpklo/hi instructions.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1250

MVT::nxv2i32 is not a legal type and so shouldn't be listed here.

[Allen]

ping?

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1250

I add type MVT::nxv2i32 because the case masked_zload_2i16_2f64 in CodeGen/AArch64/sve-masked-ldst-zext.ll.
If I don't set it Legal, then there will be an redundant and, the complete assemble is followed.

        ld1h	{ z0.d }, p0/z, [x0]
	ptrue	p1.d
	and	z0.d, z0.d, #0xffff   --- redundant ?
	ucvtf	z0.d, p1/m, z0.d

so both 64 and 128-bit vectors are supported extend ?

1250

hi, @paulwalker-arm:

do you think we need not define MVT::nxv2i32 legal, and eliminate the **and** with other method ?

[Allen]

hi, @paulwalker-arm

I find that without this refactor patch, the MVT::nxv2i32 type is also lega default, so the 'and' insn can be deleted in tryToFoldExtOfMaskedLoad.

then, if we set it expand, the condition !TLI.isLoadExtLegalOrCustom(ExtLoadType, VT, Ld->getValueType(0)) in above function tryToFoldExtOfMaskedLoad will prevent it from optimizing.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

1250

It does feel like we're missing an ISD::AND combine somewhere, but given you're just maintaining the existing behaviour I'll not worry about it for this patch. We can circle back to the MVT::nxv2i32 issue later.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

for (MVT VT : MVT::fp_scalable_vector_valuetypes()) {
         setTruncStoreAction(VT, InnerVT, Expand);
         // SVE does not have floating-point extending loads.
         setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::ZEXTLOAD, VT, InnerVT, Expand);
         setLoadExtAction(ISD::EXTLOAD, VT, InnerVT, Expand);
       }
     }
 
-    // SVE supports unpklo/hi instructions to reduce the number of loads.
-    for (auto Op : {ISD::SEXTLOAD, ISD::ZEXTLOAD, ISD::EXTLOAD}) {
-      setLoadExtAction(Op, MVT::nxv16i64, MVT::nxv16i8, Expand);
-      setLoadExtAction(Op, MVT::nxv8i64, MVT::nxv8i16, Expand);
-      setLoadExtAction(Op, MVT::nxv4i64, MVT::nxv4i32, Expand);
+    // Firstly, exclude all scalable vector extending loads/truncating stores.
+    for (MVT VT : MVT::integer_scalable_vector_valuetypes()) {
+      for (MVT InnerVT : MVT::integer_scalable_vector_valuetypes()) {
+        // TODO: truncating stores should also be exclude
+        setLoadExtAction(ISD::SEXTLOAD, VT, InnerVT, Expand);
+        setLoadExtAction(ISD::ZEXTLOAD, VT, InnerVT, Expand);
+        setLoadExtAction(ISD::EXTLOAD, VT, InnerVT, Expand);
+      }
+    }
+
+    // Then, selectively enable those which we directly support.
+    for (auto Op : {ISD::ZEXTLOAD, ISD::SEXTLOAD, ISD::EXTLOAD}) {
+      setLoadExtAction(Op, MVT::nxv2i64, MVT::nxv2i8, Legal);
+      setLoadExtAction(Op, MVT::nxv2i64, MVT::nxv2i16, Legal);
+      setLoadExtAction(Op, MVT::nxv2i64, MVT::nxv2i32, Legal);
+      setLoadExtAction(Op, MVT::nxv4i32, MVT::nxv4i8, Legal);
+      setLoadExtAction(Op, MVT::nxv2i32, MVT::nxv2i16, Legal);
+      setLoadExtAction(Op, MVT::nxv4i32, MVT::nxv4i16, Legal);
+      setLoadExtAction(Op, MVT::nxv8i16, MVT::nxv8i8, Legal);
     }
 
     // SVE supports truncating stores of 64 and 128-bit vectors
     setTruncStoreAction(MVT::v2i64, MVT::v2i8, Custom);
     setTruncStoreAction(MVT::v2i64, MVT::v2i16, Custom);
     setTruncStoreAction(MVT::v2i64, MVT::v2i32, Custom);
     setTruncStoreAction(MVT::v2i32, MVT::v2i8, Custom);
     setTruncStoreAction(MVT::v2i32, MVT::v2i16, Custom);

llvm/test/CodeGen/AArch64/sve-intrinsics-ldst-ext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -asm-verbose=1 < %s | FileCheck %s
 
 ;
-; LD1B
+; LD1SB/LD1B
 ;
 
+define <vscale x 16 x i32> @ld1b_i8_sext_i32(<vscale x 16 x i8> *%base) {
+; CHECK-LABEL: ld1b_i8_sext_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ld1sb { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    ld1sb { z1.s }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ld1sb { z2.s }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT:    ld1sb { z3.s }, p0/z, [x0, #3, mul vl]
+; CHECK-NEXT:    ret
+  %wide.load = load <vscale x 16 x i8>, <vscale x 16 x i8>* %base
+  %res = sext <vscale x 16 x i8> %wide.load to <vscale x 16 x i32>
+  ret <vscale x 16 x i32> %res
+}
+
+define <vscale x 16 x i32> @ld1b_i8_zext_i32(<vscale x 16 x i8> *%base) {
+; CHECK-LABEL: ld1b_i8_zext_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ld1b { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    ld1b { z1.s }, p0/z, [x0, #1, mul vl]
+; CHECK-NEXT:    ld1b { z2.s }, p0/z, [x0, #2, mul vl]
+; CHECK-NEXT:    ld1b { z3.s }, p0/z, [x0, #3, mul vl]
+; CHECK-NEXT:    ret
+  %wide.load = load <vscale x 16 x i8>, <vscale x 16 x i8>* %base
+  %res = zext <vscale x 16 x i8> %wide.load to <vscale x 16 x i32>
+  ret <vscale x 16 x i32> %res
+}
+
 define <vscale x 16 x i64> @ld1b_i8_sext(<vscale x 16 x i8> *%base) {
 ; CHECK-LABEL: ld1b_i8_sext:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
 ; CHECK-NEXT:    ld1sb { z0.d }, p0/z, [x0]
 ; CHECK-NEXT:    ld1sb { z1.d }, p0/z, [x0, #1, mul vl]
 ; CHECK-NEXT:    ld1sb { z2.d }, p0/z, [x0, #2, mul vl]
 ; CHECK-NEXT:    ld1sb { z3.d }, p0/z, [x0, #3, mul vl]

llvm/test/CodeGen/AArch64/sve-intrinsics-mask-ldst-ext.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve -asm-verbose=1 < %s | FileCheck %s
 
 ;
 ; LD1B
 ;
 
+define <vscale x 16 x i32> @masked_ld1b_i8_sext_i32(<vscale x 16 x i8> *%base, <vscale x 16 x i1> %mask) {
+; CHECK-LABEL: masked_ld1b_i8_sext_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ld1b { z0.b }, p0/z, [x0]
+; CHECK-NEXT:    sunpklo z1.h, z0.b
+; CHECK-NEXT:    sunpkhi z3.h, z0.b
+; CHECK-NEXT:    sunpklo z0.s, z1.h
+; CHECK-NEXT:    sunpkhi z1.s, z1.h
+; CHECK-NEXT:    sunpklo z2.s, z3.h
+; CHECK-NEXT:    sunpkhi z3.s, z3.h
+; CHECK-NEXT:    ret
+  %wide.masked.load = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* %base, i32 2, <vscale x 16 x i1> %mask, <vscale x 16 x i8> undef)
+  %res = sext <vscale x 16 x i8> %wide.masked.load to <vscale x 16 x i32>
+  ret <vscale x 16 x i32> %res
+}
+
+define <vscale x 16 x i32> @masked_ld1b_i8_zext_i32(<vscale x 16 x i8> *%base, <vscale x 16 x i1> %mask) {
+; CHECK-LABEL: masked_ld1b_i8_zext_i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ld1b { z0.b }, p0/z, [x0]
+; CHECK-NEXT:    uunpklo z1.h, z0.b
+; CHECK-NEXT:    uunpkhi z3.h, z0.b
+; CHECK-NEXT:    uunpklo z0.s, z1.h
+; CHECK-NEXT:    uunpkhi z1.s, z1.h
+; CHECK-NEXT:    uunpklo z2.s, z3.h
+; CHECK-NEXT:    uunpkhi z3.s, z3.h
+; CHECK-NEXT:    ret
+  %wide.masked.load = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* %base, i32 2, <vscale x 16 x i1> %mask, <vscale x 16 x i8> undef)
+  %res = zext <vscale x 16 x i8> %wide.masked.load to <vscale x 16 x i32>
+  ret <vscale x 16 x i32> %res
+}
+
 define <vscale x 16 x i64> @masked_ld1b_i8_sext(<vscale x 16 x i8> *%base, <vscale x 16 x i1> %mask) {
 ; CHECK-LABEL: masked_ld1b_i8_sext:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ld1b { z0.b }, p0/z, [x0]
 ; CHECK-NEXT:    sunpklo z1.h, z0.b
 ; CHECK-NEXT:    sunpkhi z0.h, z0.b
 ; CHECK-NEXT:    sunpklo z2.s, z1.h
 ; CHECK-NEXT:    sunpkhi z3.s, z1.h

llvm/test/CodeGen/AArch64/sve-masked-ldst-sext.ll

; CHECK-NEXT: ret
   %load = call <vscale x 2 x i32> @llvm.masked.load.nxv2i32(<vscale x 2 x i32> *%a, i32 1, <vscale x 2 x i1> %mask, <vscale x 2 x i32> %passthru)
   %ext = sext <vscale x 2 x i32> %load to <vscale x 2 x i64>
   ret <vscale x 2 x i64> %ext
 }
 
 ; Return type requires splitting
 define <vscale x 16 x i32> @masked_sload_nxv16i8(<vscale x 16 x i8>* %a, <vscale x 16 x i1> %mask) {
 ; CHECK-LABEL: masked_sload_nxv16i8:
-; CHECK:         punpklo p1.h, p0.b
-; CHECK-NEXT:    punpkhi p0.h, p0.b
-; CHECK-NEXT:    punpklo p2.h, p1.b
-; CHECK-NEXT:    punpkhi p1.h, p1.b
-; CHECK-NEXT:    ld1sb { z0.s }, p2/z, [x0]
-; CHECK-NEXT:    punpklo p2.h, p0.b
-; CHECK-NEXT:    punpkhi p0.h, p0.b
-; CHECK-NEXT:    ld1sb { z1.s }, p1/z, [x0, #1, mul vl]
-; CHECK-NEXT:    ld1sb { z2.s }, p2/z, [x0, #2, mul vl]
-; CHECK-NEXT:    ld1sb { z3.s }, p0/z, [x0, #3, mul vl]
+; CHECK:    ld1b { z0.b }, p0/z, [x0]
+; CHECK-NEXT:    sunpklo z1.h, z0.b
+; CHECK-NEXT:    sunpkhi z3.h, z0.b
+; CHECK-NEXT:    sunpklo z0.s, z1.h
+; CHECK-NEXT:    sunpkhi z1.s, z1.h
+; CHECK-NEXT:    sunpklo z2.s, z3.h
+; CHECK-NEXT:    sunpkhi z3.s, z3.h
 ; CHECK-NEXT:    ret
   %load = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8(<vscale x 16 x i8>* %a, i32 2, <vscale x 16 x i1> %mask, <vscale x 16 x i8> undef)
   %ext = sext <vscale x 16 x i8> %load to <vscale x 16 x i32>
   ret <vscale x 16 x i32> %ext
 }
 
 ; Masked load requires promotion
 define <vscale x 4 x double> @masked_sload_4i8_4f32(<vscale x 4 x i8>* noalias %in, <vscale x 4 x i1> %mask) {
 ; CHECK-LABEL: masked_sload_4i8_4f32:
-; CHECK:       punpkhi p2.h, p0.b
-; CHECK-NEXT:  punpklo p0.h, p0.b
-; CHECK-NEXT:  ld1sb { z1.d }, p2/z, [x0, #1, mul vl]
-; CHECK-NEXT:  ld1sb { z0.d }, p0/z, [x0]
-; CHECK-NEXT:  ptrue p1.d
-; CHECK-NEXT:  scvtf z0.d, p1/m, z0.d
-; CHECK-NEXT:  scvtf z1.d, p1/m, z1.d
-; CHECK-NEXT:  ret
+; CHECK:    ld1sb { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    sunpkhi z1.d, z0.s
+; CHECK-NEXT:    sunpklo z0.d, z0.s
+; CHECK-NEXT:    scvtf z0.d, p1/m, z0.d
+; CHECK-NEXT:    scvtf z1.d, p1/m, z1.d
+; CHECK-NEXT:    ret
   %wide.load = call <vscale x 4 x i8> @llvm.masked.load.nxv4i8(<vscale x 4 x i8>* %in, i32 2, <vscale x 4 x i1> %mask, <vscale x 4 x i8> undef)
   %sext = sext <vscale x 4 x i8> %wide.load to <vscale x 4 x i64>
   %res = sitofp <vscale x 4 x i64> %sext to <vscale x 4 x double>
   ret <vscale x 4 x double> %res
 }
 
 declare <vscale x 2 x i8> @llvm.masked.load.nxv2i8(<vscale x 2 x i8>*, i32, <vscale x 2 x i1>, <vscale x 2 x i8>)
 declare <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>*, i32, <vscale x 2 x i1>, <vscale x 2 x i16>)
 declare <vscale x 2 x i32> @llvm.masked.load.nxv2i32(<vscale x 2 x i32>*, i32, <vscale x 2 x i1>, <vscale x 2 x i32>)
 declare <vscale x 4 x i8> @llvm.masked.load.nxv4i8(<vscale x 4 x i8>*, i32, <vscale x 4 x i1>, <vscale x 4 x i8>)
 declare <vscale x 4 x i16> @llvm.masked.load.nxv4i16(<vscale x 4 x i16>*, i32, <vscale x 4 x i1>, <vscale x 4 x i16>)
 declare <vscale x 8 x i8> @llvm.masked.load.nxv8i8(<vscale x 8 x i8>*, i32, <vscale x 8 x i1>, <vscale x 8 x i8>)
 declare <vscale x 16 x i8> @llvm.masked.load.nxv16i8(<vscale x 16 x i8>*, i32, <vscale x 16 x i1>, <vscale x 16 x i8>)