fixup! [RISCV] Optimized and with atomic_load into zextload when safe.

Added pattern match for atomic_load_zext_n with replacement of getLoad with getAtomicLoad.

Author: janagor

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -15297,8 +15297,6 @@ static SDValue reduceANDOfAtomicLoad(SDNode *N,
     return SDValue();
 
   EVT LoadedVT = ALoad->getMemoryVT();
-  EVT ResultVT = N->getValueType(0);
-
   uint64_t Mask = maskTrailingOnes<uint64_t>(LoadedVT.getSizeInBits());
   uint64_t ExpectedMask = LoadedVT.getSizeInBits() == 8    ? 0xFF
                           : LoadedVT.getSizeInBits() == 16 ? 0xFFFF
@@ -15307,13 +15305,9 @@ static SDValue reduceANDOfAtomicLoad(SDNode *N,
   if (Mask != ExpectedMask)
     return SDValue();
 
-  SDLoc DL(N);
-  SDValue Chain = ALoad->getChain();
-  SDValue Ptr = ALoad->getBasePtr();
-  MachineMemOperand *MemOp = ALoad->getMemOperand();
-  SDValue ZextLoad = DAG.getExtLoad(ISD::ZEXTLOAD, DL, ResultVT, Chain, Ptr,
-                                    MemOp->getPointerInfo(), LoadedVT,
-                                    MemOp->getAlign(), MemOp->getFlags());
+  SDValue ZextLoad = DAG.getAtomicLoad(
+      ISD::ZEXTLOAD, SDLoc(N), ALoad->getMemoryVT(), N->getValueType(0),
+      ALoad->getChain(), ALoad->getBasePtr(), ALoad->getMemOperand());
   DCI.CombineTo(N, ZextLoad);
   DAG.ReplaceAllUsesOfValueWith(SDValue(N0.getNode(), 1), ZextLoad.getValue(1));
   DCI.recursivelyDeleteUnusedNodes(N0.getNode());

llvm/lib/Target/RISCV/RISCVInstrInfoA.td

@@ -167,6 +167,8 @@ class seq_cst_store<PatFrag base>
 let Predicates = [HasAtomicLdSt] in {
   def : LdPat<relaxed_load<atomic_load_asext_8>,  LB>;
   def : LdPat<relaxed_load<atomic_load_asext_16>, LH>;
+  def : LdPat<relaxed_load<atomic_load_zext_8>,  LBU>;
+  def : LdPat<relaxed_load<atomic_load_zext_16>, LHU>;
 
   def : StPat<relaxed_store<atomic_store_8>,  SB, GPR, XLenVT>;
   def : StPat<relaxed_store<atomic_store_16>, SH, GPR, XLenVT>;
@@ -179,6 +181,7 @@ let Predicates = [HasAtomicLdSt, IsRV32] in {
 
 let Predicates = [HasAtomicLdSt, IsRV64] in {
   def : LdPat<relaxed_load<atomic_load_asext_32>, LW>;
+  def : LdPat<relaxed_load<atomic_load_zext_32>, LWU>;
   def : LdPat<relaxed_load<atomic_load_64>, LD, i64>;
   def : StPat<relaxed_store<atomic_store_64>, SD, GPR, i64>;
 }