[DAG] Allow folding AND of anyext masked_load with >1 user to zext version

[benmxwl-arm]

Summary

This now allows folding an AND of a anyext masked_load to a zext_masked_load even if the masked load has multiple users.
Doing is eliminates some redundant ANDs/MOVs for certain AArch64 SVE code.

I'm not sure if there's any cases where doing this could negatively the other users of the masked_load.
Looking at other optimizations of masked loads, most don't apply if the load is used more than once, so it doesn't look like this would interfere.

[benmxwl-arm]

Remove accidentally committed file

[benmxwl-arm]

Replace the loads in the same way tryToFoldExtOfLoad() handles it. This avoids some errors in the DAG.

[benmxwl-arm]

Corrected isTruncateFree() check for existing users.

[peterwaller-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12121

V is N0, so do you need to close over N0?

12131

Is it necessary to bail out for these cases, or would they still be handled by the isTruncFree logic? If you want to have the FIXME, it may be worth referencing ExtendUsesToFormExtLoad.

[benmxwl-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

12121

I also close over VT and TLI, so I guess I could just use the closed N0 too.

12131

I'm not sure it's covered by the isTruncFree logic. The SETCC case is handled for the other loads by keeping a list of the SETCC nodes, then extending their other operands to match the new extended load.

(This is only done for sign-extension, otherwise the transform load transform is not applied).

When there's live in/out nodes the number of SETCCs is used to decide if the transform is worth it.

[benmxwl-arm]

Updated to a simpler solution for the same problem.

This now allows folding an AND of a anyext masked_load to a zext_masked_load even if the masked load has multiple users.
Doing is eliminates some redundant ANDs/MOVs for certain AArch64 SVE code.

I'm not sure if there's any cases where doing this could negatively the other users of the masked_load.
Looking at other optimizations of masked loads, most don't apply if the load is used more than once, so it doesn't look like this would interfere.

[benmxwl-arm]

Marking Peter's comments as done as they no longer apply to this version.

[c-rhodes]

this looks ok to me, but it would be good to get confirmation it's ok to remove the hasOneUse check on the masked load, @samtebbs @RKSimon any reason we might not want to do this?

[samtebbs]

From my reading an anyext can't guarantee what is put in those extended bits so changing the anyext to a zext seems semantically acceptable to me. If it isn't correct then we'll get some bug reports and it can be fixed!

[c-rhodes]

In D137844#3933258, @samtebbs wrote:

From my reading an anyext can't guarantee what is put in those extended bits so changing the anyext to a zext seems semantically acceptable to me. If it isn't correct then we'll get some bug reports and it can be fixed!

Great, sounds good. thanks for taking a look. LGTM.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (VT.isVector()) {
 
     // fold (and x, -1) -> x, vector edition
     if (ISD::isConstantSplatVectorAllOnes(N1.getNode()))
       return N0;
 
     // fold (and (masked_load) (splat_vec (x, ...))) to zext_masked_load
     auto *MLoad = dyn_cast<MaskedLoadSDNode>(N0);
     ConstantSDNode *Splat = isConstOrConstSplat(N1, true, true);
-    if (MLoad && MLoad->getExtensionType() == ISD::EXTLOAD && N0.hasOneUse() &&
-        Splat && N1.hasOneUse()) {
+    if (MLoad && MLoad->getExtensionType() == ISD::EXTLOAD && Splat &&
+        N1.hasOneUse()) {
       EVT LoadVT = MLoad->getMemoryVT();
       EVT ExtVT = VT;
       if (TLI.isLoadExtLegal(ISD::ZEXTLOAD, ExtVT, LoadVT)) {
         // For this AND to be a zero extension of the masked load the elements
         // of the BuildVec must mask the bottom bits of the extended element
         // type
         uint64_t ElementSize =
             LoadVT.getVectorElementType().getScalarSizeInBits();
         if (Splat->getAPIntValue().isMask(ElementSize)) {
-          return DAG.getMaskedLoad(
+          auto NewLoad = DAG.getMaskedLoad(
               ExtVT, SDLoc(N), MLoad->getChain(), MLoad->getBasePtr(),
               MLoad->getOffset(), MLoad->getMask(), MLoad->getPassThru(),
               LoadVT, MLoad->getMemOperand(), MLoad->getAddressingMode(),
               ISD::ZEXTLOAD, MLoad->isExpandingLoad());
+          bool LoadHasOtherUsers = !N0.hasOneUse();
+          CombineTo(N, NewLoad);
+          if (LoadHasOtherUsers)
+            CombineTo(MLoad, NewLoad.getValue(0), NewLoad.getValue(1));
+          return SDValue(N, 0);
         }
       }
     }
   }
 
   // fold (and x, -1) -> x
   if (isAllOnesConstant(N1))
     return N0;
static SDValue tryToFoldExtOfMaskedLoad(SelectionDAG &DAG,
 
   if (!TLI.isLoadExtLegalOrCustom(ExtLoadType, VT, Ld->getValueType(0)))
     return SDValue();
 
   if (!TLI.isVectorLoadExtDesirable(SDValue(N, 0)))
     return SDValue();
 
   SDLoc dl(Ld);
   SDValue PassThru = DAG.getNode(ExtOpc, dl, VT, Ld->getPassThru());
   SDValue NewLoad = DAG.getMaskedLoad(
       VT, dl, Ld->getChain(), Ld->getBasePtr(), Ld->getOffset(), Ld->getMask(),
       PassThru, Ld->getMemoryVT(), Ld->getMemOperand(), Ld->getAddressingMode(),
       ExtLoadType, Ld->isExpandingLoad());
   DAG.ReplaceAllUsesOfValueWith(SDValue(Ld, 1), SDValue(NewLoad.getNode(), 1));
   return NewLoad;
 }
 
 static SDValue foldExtendedSignBitTest(SDNode *N, SelectionDAG &DAG,
                                        bool LegalOperations) {
   assert((N->getOpcode() == ISD::SIGN_EXTEND ||
           N->getOpcode() == ISD::ZERO_EXTEND) && "Expected sext or zext");
 
   SDValue SetCC = N->getOperand(0);
   if (LegalOperations || SetCC.getOpcode() != ISD::SETCC ||
       !SetCC.hasOneUse() || SetCC.getValueType() != MVT::i1)
     return SDValue();
 

llvm/test/CodeGen/AArch64/sve-load-compare-store.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
 
 define void @sve_load_compare_store(ptr noalias nocapture noundef readonly %a, ptr noalias nocapture noundef %b) {
 ; CHECK-LABEL: sve_load_compare_store:
 ; CHECK:       // %bb.0: // %entry
 ; CHECK-NEXT:    ptrue p0.b
 ; CHECK-NEXT:    ld1h { z0.s }, p0/z, [x0]
-; CHECK-NEXT:    mov z1.d, z0.d
-; CHECK-NEXT:    and z1.s, z1.s, #0xffff
-; CHECK-NEXT:    cmphs p0.s, p0/z, z1.s, #0
+; CHECK-NEXT:    cmphs p0.s, p0/z, z0.s, #0
 ; CHECK-NEXT:    st1b { z0.s }, p0, [x1]
 ; CHECK-NEXT:    ret
 entry:
   %0 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.ptrue.nxv16i1(i32 31)
   %1 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv4i1(<vscale x 16 x i1> %0)
   %2 = tail call <vscale x 4 x i16> @llvm.masked.load.nxv4i16.p0(ptr %a, i32 1, <vscale x 4 x i1> %1, <vscale x 4 x i16> zeroinitializer)
   %3 = zext <vscale x 4 x i16> %2 to <vscale x 4 x i32>
   %4 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.cmphs.nxv4i32(<vscale x 4 x i1> %1, <vscale x 4 x i32> %3, <vscale x 4 x i32> zeroinitializer)

llvm/test/CodeGen/Thumb2/mve-masked-load.ll

 ; CHECK-LE-NEXT:    push {r7, lr}
 ; CHECK-LE-NEXT:    .vsave {d8, d9}
 ; CHECK-LE-NEXT:    vpush {d8, d9}
 ; CHECK-LE-NEXT:    vpt.s32 gt, q0, zr
 ; CHECK-LE-NEXT:    vldrht.u32 q4, [r0]
 ; CHECK-LE-NEXT:    vmov r0, r1, d8
 ; CHECK-LE-NEXT:    vmov r2, r3, d9
 ; CHECK-LE-NEXT:    bl foo
-; CHECK-LE-NEXT:    vmovlb.u16 q0, q4
+; CHECK-LE-NEXT:    vmov q0, q4
 ; CHECK-LE-NEXT:    vpop {d8, d9}
 ; CHECK-LE-NEXT:    pop {r7, pc}
 ;
 ; CHECK-BE-LABEL: multi_user_zext:
 ; CHECK-BE:       @ %bb.0: @ %entry
 ; CHECK-BE-NEXT:    .save {r7, lr}
 ; CHECK-BE-NEXT:    push {r7, lr}
 ; CHECK-BE-NEXT:    .vsave {d8, d9}
 ; CHECK-BE-NEXT:    vpush {d8, d9}
 ; CHECK-BE-NEXT:    vrev64.32 q1, q0
 ; CHECK-BE-NEXT:    vpt.s32 gt, q1, zr
-; CHECK-BE-NEXT:    vldrht.u32 q4, [r0]
-; CHECK-BE-NEXT:    vrev64.32 q0, q4
-; CHECK-BE-NEXT:    vmov r1, r0, d0
-; CHECK-BE-NEXT:    vmov r3, r2, d1
+; CHECK-BE-NEXT:    vldrht.u32 q0, [r0]
+; CHECK-BE-NEXT:    vrev64.32 q4, q0
+; CHECK-BE-NEXT:    vmov r1, r0, d8
+; CHECK-BE-NEXT:    vmov r3, r2, d9
 ; CHECK-BE-NEXT:    bl foo
-; CHECK-BE-NEXT:    vmovlb.u16 q1, q4
-; CHECK-BE-NEXT:    vrev64.32 q0, q1
+; CHECK-BE-NEXT:    vmov q0, q4
 ; CHECK-BE-NEXT:    vpop {d8, d9}
 ; CHECK-BE-NEXT:    pop {r7, pc}
 entry:
   %c = icmp sgt <4 x i32> %a, zeroinitializer
   %l = call <4 x i16> @llvm.masked.load.v4i16.p0v4i16(<4 x i16>* %dest, i32 2, <4 x i1> %c, <4 x i16> zeroinitializer)
   call void @foo(<4 x i16> %l)
   %ext = zext <4 x i16> %l to <4 x i32>
   ret <4 x i32> %ext