[llvm][SVE] Remove redundant and when comparing against extending load

[DavidTruby]

Summary

When determining if an and should be merged into an extending load
the constant argument to the and is currently not checked if the
argument requires truncation. This prevents the combine happening when
the vector width is half the normal available vector width for SVE VLA
vectors.

[DavidTruby]

Since there were no tests triggering this before, here's a compiler explorer link showing what codegen looks like at the moment for comparison: https://godbolt.org/z/3f8Wh5c5x

I've added more tests than just the one that gets triggered here as more of these have the same issue on llvm 14 so that should also be defended against regressions.

[c-rhodes]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if ((N0.getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
       (N0.getOpcode() == ISD::LOAD && N0.getResNo() == 0)) {
     LoadSDNode *Load = cast<LoadSDNode>( (N0.getOpcode() == ISD::LOAD) ?
                                          N0 : N0.getOperand(0) );
 
     // Get the constant (if applicable) the zero'th operand is being ANDed with.
     // This can be a pure constant or a vector splat, in which case we treat the
     // vector as a scalar and use the splat value.
     APInt Constant = APInt::getZero(1);
-    if (const ConstantSDNode *C = isConstOrConstSplat(N1)) {
+    if (const ConstantSDNode *C = isConstOrConstSplat(
+            N1, /*AllowUndef=*/false, /*AllowTruncation=*/true)) {
       Constant = C->getAPIntValue();
     } else if (BuildVectorSDNode *Vector = dyn_cast<BuildVectorSDNode>(N1)) {
       APInt SplatValue, SplatUndef;
       unsigned SplatBitSize;
       bool HasAnyUndefs;
       bool IsSplat = Vector->isConstantSplat(SplatValue, SplatUndef,
                                              SplatBitSize, HasAnyUndefs);
       if (IsSplat) {

llvm/test/CodeGen/AArch64/sve-extload-icmp.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc < %s | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define <vscale x 8 x i8> @extload_icmp_nxv8i8(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv8i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    ld1b { z0.h }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 8 x i8>, ptr %in
+  %cmp = icmp eq <vscale x 8 x i8> %ld, zeroinitializer
+  %ex = zext <vscale x 8 x i1> %cmp to <vscale x 8 x i8>
+  ret <vscale x 8 x i8> %ex
+}
+
+define <vscale x 16 x i8> @extload_icmp_nxv16i8(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv16i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    ld1b { z0.b }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.b, p0/m, z0.b
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 16 x i8>, ptr %in
+  %cmp = icmp eq <vscale x 16 x i8> %ld, zeroinitializer
+  %ex = zext <vscale x 16 x i1> %cmp to <vscale x 16 x i8>
+  ret <vscale x 16 x i8> %ex
+}
+
+define <vscale x 4 x i16> @extload_icmp_nxv4i16(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv4i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ld1h { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 4 x i16>, ptr %in
+  %cmp = icmp eq <vscale x 4 x i16> %ld, zeroinitializer
+  %ex = zext <vscale x 4 x i1> %cmp to <vscale x 4 x i16>
+  ret <vscale x 4 x i16> %ex
+}
+
+define <vscale x 8 x i16> @extload_icmp_nxv8i16(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv8i16:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.h
+; CHECK-NEXT:    ld1h { z0.h }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.h, p0/m, z0.h
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 8 x i16>, ptr %in
+  %cmp = icmp eq <vscale x 8 x i16> %ld, zeroinitializer
+  %ex = zext <vscale x 8 x i1> %cmp to <vscale x 8 x i16>
+  ret <vscale x 8 x i16> %ex
+}
+
+define <vscale x 2 x i32> @extload_icmp_nxv2i32(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv2i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    ld1w { z0.d }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.d, p0/m, z0.d
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 2 x i32>, ptr %in
+  %cmp = icmp eq <vscale x 2 x i32> %ld, zeroinitializer
+  %ex = zext <vscale x 2 x i1> %cmp to <vscale x 2 x i32>
+  ret <vscale x 2 x i32> %ex
+}
+
+define <vscale x 4 x i32> @extload_icmp_nxv4i32(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv4i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    ld1w { z0.s }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.s, p0/m, z0.s
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 4 x i32>, ptr %in
+  %cmp = icmp eq <vscale x 4 x i32> %ld, zeroinitializer
+  %ex = zext <vscale x 4 x i1> %cmp to <vscale x 4 x i32>
+  ret <vscale x 4 x i32> %ex
+}
+
+define <vscale x 2 x i64> @extload_icmp_nxv2i64(ptr %in) #0 {
+; CHECK-LABEL: extload_icmp_nxv2i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    ld1d { z0.d }, p0/z, [x0]
+; CHECK-NEXT:    cnot z0.d, p0/m, z0.d
+; CHECK-NEXT:    ret
+  %ld = load <vscale x 2 x i64>, ptr %in
+  %cmp = icmp eq <vscale x 2 x i64> %ld, zeroinitializer
+  %ex = zext <vscale x 2 x i1> %cmp to <vscale x 2 x i64>
+  ret <vscale x 2 x i64> %ex
+}
+
+
+
+attributes #0 = { "target-features"="+sve" }