[AArch64][InstCombine] Fold MLOAD and zero extensions into MLOAD

[Allen]

Summary

Accord the discussion in D122281, we missing an ISD::AND combine for MLOAD
because it relies on BuildVectorSDNode is fails for scalable vectors.
This patch is intend to handle that, so we can circle back the type MVT::nxv2i32

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14422–14423

The masked load's MemVT is not sufficient as you must also consider the load's extension type.

1. As it stand this transform is only safe for zero-extending masked loads...
2. for an any existing masked load the extension can be changed to zero-extending...
3. for a single use sign extending load the extension can be changed to zero-extending.

How far down this path you go depends on the exact use case you need. If it's just (1) then I think this can be a common DAGCombine. For (2) and (3) it'll depend on whether there are sufficient TLI hooks to ensure you don't create a version a target cannot lower. That said, given you require the input to be an extending masked load perhaps that is not a problem.

One final issue is the handling of the masked load's pass through value. With the AND in place this will be zero-extended but once removed that stops and thus that will need to be considered as well. My guess is that for today's use case you want to restrict the combine to instances where the pass through value is either undef or zero.

[Allen]

hi @paulwalker-arm

With some debug with function DAGCombiner::visitAND, I find there is another issue need to check:

it use TLI.isLoadExtLegal(ISD::ZEXTLOAD, ExtVT, LoadVT) to check whether the following transform can be done, do we need similar check in the AArch64 backend? If yes, I think we must set MVT::nxv2i32 legal.
fold (and (masked_load) (build_vec (x, ...))) to zext_masked_load

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14422–14423

Thanks @paulwalker-arm for detail comment.
I has a question: Both 2) any extension and 3) sign extension can be changed to zero-extending, so why we need check the load's extension type ?

[paulwalker-arm]

In D122703#3418244, @Allen wrote:

hi @paulwalker-arm

With some debug with function DAGCombiner::visitAND, I find there is another issue need to check:

't s
it use TLI.isLoadExtLegal(ISD::ZEXTLOAD, ExtVT, LoadVT) to check whether the following transform can be done, do we need similar check in the AArch64 backend? If yes, I think we must set MVT::nxv2i32 legal.

Not sure I understand what you mean here. ExtVT will be the legal type, most likely MVT::nxv2i64, whereas MVT::nxv2i32 will be the MemVT that doesn't need to be legal, so I don't see an issue? Looking at DAGCombiner::visitAND it looks like it already does a related transformation but because it relies on BuildVectorSDNode is fails for scalable vectors. To me it looks like we can make this common code more portable and support all vector types.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

14422–14423

The need to check the extension type depends on which of the options you go with. If you only implement (1) then you need to check the extension type is explicitly zero-extending. Options (1) & (2) require you to skip sign-extending cases and for (3) you'll still need a use count check for the sign-extending case.

[Allen]

rewrite to support splat_vec in DAGCombiner::visitAND

[Allen]

In D122703#3419784, @paulwalker-arm wrote:

In D122703#3418244, @Allen wrote:

hi @paulwalker-arm

With some debug with function DAGCombiner::visitAND, I find there is another issue need to check:

't s
it use TLI.isLoadExtLegal(ISD::ZEXTLOAD, ExtVT, LoadVT) to check whether the following transform can be done, do we need similar check in the AArch64 backend? If yes, I think we must set MVT::nxv2i32 legal.

Not sure I understand what you mean here. ExtVT will be the legal type, most likely MVT::nxv2i64, whereas MVT::nxv2i32 will be the MemVT that doesn't need to be legal, so I don't see an issue? Looking at DAGCombiner::visitAND it looks like it already does a related transformation but because it relies on BuildVectorSDNode is fails for scalable vectors. To me it looks like we can make this common code more portable and support all vector types.

I rewrite the MR with above review with support splat_vec in DAGCombiner::visitAND, thanks @paulwalker-arm very much for detail advice.

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

6055–6056

If you follow the advice below I'd also remove the build_vec comment as the splat_vec one is meaning full enough in itself.

6058

You could simplify the code by extract the splat here, i.e.

ConstantSDNode *Splat = isConstOrConstSplat(N1, true, true)

That would remove the separate buildvector and splatvector checks and also remove a level of indentation in the TLI.isLoadExtLegal block.

[Allen]

update according the review

[paulwalker-arm]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

6056

Sorry I think I wasn't clear enough with my previous comment. I meant that it is not worth having both the original build_vec and your new splat_vec comment and so it would be better to just keep the splat_vec one (i.e. fold (and (masked_load) (splat_vec (x, ...))) to zext_masked_load)

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (VT.isVector()) {
     if (ISD::isConstantSplatVectorAllZeros(N1.getNode()))
       // do not return N1, because undef node may exist in N1
       return DAG.getConstant(APInt::getZero(N1.getScalarValueSizeInBits()),
                              SDLoc(N), N1.getValueType());
 
     // fold (and x, -1) -> x, vector edition
     if (ISD::isConstantSplatVectorAllOnes(N1.getNode()))
       return N0;
 
-    // fold (and (masked_load) (build_vec (x, ...))) to zext_masked_load
     auto *MLoad = dyn_cast<MaskedLoadSDNode>(N0);
-    auto *BVec = dyn_cast<BuildVectorSDNode>(N1);
-    if (MLoad && BVec && MLoad->getExtensionType() == ISD::EXTLOAD &&
-        N0.hasOneUse() && N1.hasOneUse()) {
+    ConstantSDNode *Splat = isConstOrConstSplat(N1, true, true);
+    if (MLoad && MLoad->getExtensionType() == ISD::EXTLOAD && N0.hasOneUse() &&
+        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
-        if (ConstantSDNode *Splat = BVec->getConstantSplatNode()) {
-          uint64_t ElementSize =
-              LoadVT.getVectorElementType().getScalarSizeInBits();
-          if (Splat->getAPIntValue().isMask(ElementSize)) {
-            return DAG.getMaskedLoad(
-                ExtVT, SDLoc(N), MLoad->getChain(), MLoad->getBasePtr(),
-                MLoad->getOffset(), MLoad->getMask(), MLoad->getPassThru(),
-                LoadVT, MLoad->getMemOperand(), MLoad->getAddressingMode(),
-                ISD::ZEXTLOAD, MLoad->isExpandingLoad());
-          }
-        }
-      }
-    }
-  }
 
   // fold (and x, -1) -> x
   if (isAllOnesConstant(N1))
     return N0;
 
   // if (and x, c) is known to be zero, return 0
   unsigned BitWidth = VT.getScalarSizeInBits();
   ConstantSDNode *N1C = isConstOrConstSplat(N1);

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Subtarget->hasSVE()) {
     }
 
     // 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);
case AArch64ISD::LDFF1_MERGE_ZERO:
     break;
   case AArch64ISD::GLD1_MERGE_ZERO:
   case AArch64ISD::GLD1_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_MERGE_ZERO:
   case AArch64ISD::GLD1_UXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLD1_IMM_MERGE_ZERO:
   case AArch64ISD::GLDFF1_MERGE_ZERO:
   case AArch64ISD::GLDFF1_SCALED_MERGE_ZERO:
   case AArch64ISD::GLDFF1_SXTW_MERGE_ZERO:
   case AArch64ISD::GLDFF1_SXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLDFF1_UXTW_MERGE_ZERO:
   case AArch64ISD::GLDFF1_UXTW_SCALED_MERGE_ZERO:
   case AArch64ISD::GLDFF1_IMM_MERGE_ZERO:
   case AArch64ISD::GLDNT1_MERGE_ZERO:
     MemVT = cast<VTSDNode>(Src->getOperand(4))->getVT();
     break;

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

; CHECK-NEXT:    ret
   ret <vscale x 8 x i64> %ext
 }
 
 ; Masked load requires promotion
 define <vscale x 2 x double> @masked_zload_2i16_2f64(<vscale x 2 x i16>* noalias %in, <vscale x 2 x i1> %mask) {
 ; CHECK-LABEL: masked_zload_2i16_2f64:
 ; CHECK:       ld1h { z0.d }, p0/z, [x0]
 ; CHECK-NEXT:  ptrue p0.d
-; CHECK-NEXT:  ucvtf z0.d, p0/m, z0.s
+; CHECK-NEXT:  ucvtf z0.d, p0/m, z0.d
 ; CHECK-NEXT:  ret
   %wide.load = call <vscale x 2 x i16> @llvm.masked.load.nxv2i16(<vscale x 2 x i16>* %in, i32 2, <vscale x 2 x i1> %mask, <vscale x 2 x i16> undef)
   %zext = zext <vscale x 2 x i16> %wide.load to <vscale x 2 x i32>
   %res = uitofp <vscale x 2 x i32> %zext to <vscale x 2 x double>
   ret <vscale x 2 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 8 x i16> @llvm.masked.load.nxv8i16(<vscale x 8 x i16>*, i32, <vscale x 8 x i1>, <vscale x 8 x i16>)