[DAGCombine] Refactor ReduceLoadWidth

[samparker]

Summary

visitAND attempts to narrow the width of extending loads that are then masked off. ReduceLoadWidth already exists for a similar purpose and handles shifts, so I've moved the code to handle AND nodes there.

[fhahn]

Removing the unneeded LoadedVT seems like a straight-forward NFC and I think it would be worth to submit it for review as an isolated change, as it does not seem related to moving some code to ReduceLoadWidth.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

7969 ↗

(On Diff #121466)

If we support extension using AND, I think we should also update the comment here.

7999 ↗

(On Diff #121466)

I am probably missing something, but I am not sure if I see the real benefit of moving this code here. It seems the visitAND was the right place to do it, without requiring additional checks. However, it seems like there is potential to simplify the code in visitAND

[samparker]

Hi Florian,

Thanks for the comments, I will revert the LoadedVT change.

cheers,
sam

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

7999 ↗

(On Diff #121466)

Both visitAND and ReduceLoadWidth contained essentially the same code to generate the new load, so it seems reasonable to me to contain this logic within the sensibly named function. In D39604 I also wish to directly use this function to combine ANDs and hopefully it can more helpful when anyone else wants to do similar.

[samparker]

Removed the unrelated API change and updated the description of ReduceLoadWidth.

[samparker]

Hi Florian,

Do you have any other suggestions?

cheers

[fhahn]

Looks like a reasonable refactor to me helping with D39604. Adding Simon & Eli as reviewers as I think they know the area better then myself.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

7999 ↗

(On Diff #121466)

Ah yes, there is some shared code to generated the new load, I missed initially! ReduceLoadWidth also contains a few more legality checks.

[efriedma]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8026 ↗

(On Diff #121703)

Could this block of code be simplified? As long as you set ExtType and ExtVT correctly, I think the rest of the checks already exist in some form, by analogy to the transforms for ISD::SIGN_EXTEND_IN_REG.

[samparker]

Removed most of the checks performed on the load, leaving just the SEXT check.

[samparker]

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8026 ↗

(On Diff #121703)

Thanks Eli, indeed most of the checks that performed again.

[fhahn]

Thanks Sam, looks good to me. Please wait a day or two with committing, to give others a chance to raise additional concerns.

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8034 ↗

(On Diff #123311)

IIRC your first version of the patch removed to unneeded LoadedVT argument. I think this would be a worthwhile simplification to get in as NFC independently.

[RKSimon]

Very minor comment

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

8031 ↗

(On Diff #123311)

use bool here - we typically only use auto for casts

[samparker]

@fhahn @RKSimon @eli.friedman I made some small changes before committing, could someone please check them out? My development branch only had x86, Arm and Aarch64 targets built and so when I went to run my pre-commit tests, some were failing for AMDGPU and NVPTX. The issues were that even though some checks are duplicated, I have an early return in the new block which still requires those checks. Also I had missed how any_extend nodes are combined, so I re-added some code that was lost during the refactor.

[efriedma]

I'm still kind of concerned that this is messier than it needs to be. I guess I'll try experimenting a bit.

[efriedma]

See D40667.

llvm/trunk/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitAND(SDNode *N) {
   }
 
   // fold (and (load x), 255) -> (zextload x, i8)
   // fold (and (extload x, i16), 255) -> (zextload x, i8)
   // fold (and (any_ext (extload x, i16)), 255) -> (zextload x, i8)
   if (!VT.isVector() && N1C && (N0.getOpcode() == ISD::LOAD ||
                                 (N0.getOpcode() == ISD::ANY_EXTEND &&
                                  N0.getOperand(0).getOpcode() == ISD::LOAD))) {
-    bool HasAnyExt = N0.getOpcode() == ISD::ANY_EXTEND;
-    LoadSDNode *LN0 = HasAnyExt
-      ? cast<LoadSDNode>(N0.getOperand(0))
-      : cast<LoadSDNode>(N0);
-    if (LN0->getExtensionType() != ISD::SEXTLOAD &&
-        LN0->isUnindexed() && N0.hasOneUse() && SDValue(LN0, 0).hasOneUse()) {
-      auto NarrowLoad = false;
-      EVT LoadResultTy = HasAnyExt ? LN0->getValueType(0) : VT;
-      EVT ExtVT, LoadedVT;
-      if (isAndLoadExtLoad(N1C, LN0, LoadResultTy, ExtVT, LoadedVT,
-                           NarrowLoad)) {
-        if (!NarrowLoad) {
-          SDValue NewLoad =
-            DAG.getExtLoad(ISD::ZEXTLOAD, SDLoc(LN0), LoadResultTy,
-                           LN0->getChain(), LN0->getBasePtr(), ExtVT,
-                           LN0->getMemOperand());
-          AddToWorklist(N);
-          CombineTo(LN0, NewLoad, NewLoad.getValue(1));
-          return SDValue(N, 0);   // Return N so it doesn't get rechecked!
-        } else {
-          EVT PtrType = LN0->getOperand(1).getValueType();
+    if (SDValue Res = ReduceLoadWidth(N)) {
+      LoadSDNode *LN0 = N0->getOpcode() == ISD::ANY_EXTEND
+        ? cast<LoadSDNode>(N0.getOperand(0)) : cast<LoadSDNode>(N0);
 
-          unsigned Alignment = LN0->getAlignment();
-          SDValue NewPtr = LN0->getBasePtr();
-
-          // For big endian targets, we need to add an offset to the pointer
-          // to load the correct bytes.  For little endian systems, we merely
-          // need to read fewer bytes from the same pointer.
-          if (DAG.getDataLayout().isBigEndian()) {
-            unsigned LVTStoreBytes = LoadedVT.getStoreSize();
-            unsigned EVTStoreBytes = ExtVT.getStoreSize();
-            unsigned PtrOff = LVTStoreBytes - EVTStoreBytes;
-            SDLoc DL(LN0);
-            NewPtr = DAG.getNode(ISD::ADD, DL, PtrType,
-                                 NewPtr, DAG.getConstant(PtrOff, DL, PtrType));
-            Alignment = MinAlign(Alignment, PtrOff);
-          }
-
-          AddToWorklist(NewPtr.getNode());
-
-          SDValue Load = DAG.getExtLoad(
-              ISD::ZEXTLOAD, SDLoc(LN0), LoadResultTy, LN0->getChain(), NewPtr,
-              LN0->getPointerInfo(), ExtVT, Alignment,
-              LN0->getMemOperand()->getFlags(), LN0->getAAInfo());
-          AddToWorklist(N);
-          CombineTo(LN0, Load, Load.getValue(1));
-          return SDValue(N, 0);   // Return N so it doesn't get rechecked!
-        }
-      }
+      CombineTo(LN0, Res, Res.getValue(1));
+      return SDValue(N, 0);
     }
   }
 
   if (SDValue Combined = visitANDLike(N0, N1, N))
     return Combined;
 
   // Simplify: (and (op x...), (op y...))  -> (op (and x, y))
   if (N0.getOpcode() == N1.getOpcode())
SDValue DAGCombiner::visitAssertExt(SDNode *N) {
   }
 
   return SDValue();
 }
 
 /// If the result of a wider load is shifted to right of N  bits and then
 /// truncated to a narrower type and where N is a multiple of number of bits of
 /// the narrower type, transform it to a narrower load from address + N / num of
-/// bits of new type. If the result is to be extended, also fold the extension
-/// to form a extending load.
+/// bits of new type. Also narrow the load if the result is masked with an AND
+/// to effectively produce a smaller type. If the result is to be extended, also
+/// fold the extension to form a extending load.
 SDValue DAGCombiner::ReduceLoadWidth(SDNode *N) {
   unsigned Opc = N->getOpcode();
 
   ISD::LoadExtType ExtType = ISD::NON_EXTLOAD;
   SDValue N0 = N->getOperand(0);
   EVT VT = N->getValueType(0);
   EVT ExtVT = VT;
 
if (Opc == ISD::SIGN_EXTEND_INREG) {
 
     uint64_t ShiftAmt = N01->getZExtValue();
     uint64_t MemoryWidth = LN0->getMemoryVT().getSizeInBits();
     if (LN0->getExtensionType() != ISD::SEXTLOAD && MemoryWidth > ShiftAmt)
       ExtVT = EVT::getIntegerVT(*DAG.getContext(), MemoryWidth - ShiftAmt);
     else
       ExtVT = EVT::getIntegerVT(*DAG.getContext(),
                                 VT.getSizeInBits() - ShiftAmt);
+  } else if (Opc == ISD::AND) {
+    bool HasAnyExt = N0.getOpcode() == ISD::ANY_EXTEND;
+    LoadSDNode *LN0 =
+      HasAnyExt ? cast<LoadSDNode>(N0.getOperand(0)) : cast<LoadSDNode>(N0);
+
+    if (LN0->getExtensionType() == ISD::SEXTLOAD ||
+        !LN0->isUnindexed() || !N0.hasOneUse() || !SDValue(LN0, 0).hasOneUse())
+      return SDValue();
+
+    auto N1C = dyn_cast<ConstantSDNode>(N->getOperand(1));
+    if (!N1C)
+      return SDValue();
+
+    EVT LoadedVT;
+    bool NarrowLoad = false;
+    ExtType = ISD::ZEXTLOAD;
+    VT = HasAnyExt ? LN0->getValueType(0) : VT;
+    if (!isAndLoadExtLoad(N1C, LN0, VT, ExtVT, LoadedVT, NarrowLoad))
+      return SDValue();
+
+    if (!NarrowLoad)
+      return DAG.getExtLoad(ISD::ZEXTLOAD, SDLoc(LN0), VT,
+                            LN0->getChain(), LN0->getBasePtr(), ExtVT,
+                            LN0->getMemOperand());
   }
   if (LegalOperations && !TLI.isLoadExtLegal(ExtType, VT, ExtVT))
     return SDValue();
 
   unsigned EVTBits = ExtVT.getSizeInBits();
 
   // Do not generate loads of non-round integer types since these can
   // be expensive (and would be wrong if the type is not byte sized).