[llvm] [GISel] Combine vector load followed by an extractelement (PR #72670)

[Matt Arsenault via llvm-commits]

================
@@ -1165,6 +1165,140 @@ bool CombinerHelper::findPreIndexCandidate(GLoadStore &LdSt, Register &Addr,
   return RealUse;
 }
 
+bool CombinerHelper::matchCombineExtractedVectorLoad(MachineInstr &MI) {
+  assert(MI.getOpcode() == TargetOpcode::G_EXTRACT_VECTOR_ELT);
+
+  // Check if there is a load that defines the vector being extracted from.
+  MachineInstr *LoadMI =
+      getOpcodeDef(TargetOpcode::G_LOAD, MI.getOperand(1).getReg(), MRI);
+  if (!LoadMI)
+    return false;
+
+  Register Vector = MI.getOperand(1).getReg();
+  LLT VecEltVT = MRI.getType(Vector).getElementType();
+  LLT ResultVT = MRI.getType(MI.getOperand(0).getReg());
+
+  // Do not combine when result type and vector element type are not the same.
+  if (ResultVT != VecEltVT)
+    return false;
+
+  // Checking whether we should reduce the load width.
+  if (VecEltVT.isVector() || !MRI.hasOneUse(Vector))
+    return false;
+
+  GLoadStore *GLoadMI = cast<GLoadStore>(LoadMI);
+
+  // Check if the defining load is simple.
+  if (!GLoadMI->isSimple())
+    return false;
+
+  // If the vector element type is not a multiple of a byte then we are unable
+  // to correctly compute an address to load only the extracted element as a
+  // scalar.
+  if (!VecEltVT.isByteSized())
+    return false;
+
+  // Check if the new load that we are going to create is legal
+  // if we are in the post-legalization phase.
+  MachineMemOperand MMO = GLoadMI->getMMO();
+  MachinePointerInfo PtrInfo;
+
+  Register Index = MI.getOperand(2).getReg();
+
+  // Finding the appropriate PtrInfo if offset is a known constant.
+  // This is required to create the memory operand for the narrowed load.
+  // This machine memory operand object helps us infer about legality
+  // before we proceed to combine the instruction.
+  if (MRI.getVRegDef(Index)->getOpcode() == TargetOpcode::G_CONSTANT) {
+    MachineInstr *ConstMI = MRI.getVRegDef(Index);
+    const ConstantInt *CVal = ConstMI->getOperand(1).getCImm();
+    int Elt = CVal->getZExtValue();
+    // FIXME: should be (ABI size)*Elt.
+    unsigned PtrOff = VecEltVT.getSizeInBits() * Elt / 8;
+    PtrInfo = MMO.getPointerInfo().getWithOffset(PtrOff);
+  } else {
+    // Discard the pointer info except the address space because the memory
+    // operand can't represent this new access since the offset is variable.
+    PtrInfo = MachinePointerInfo(MMO.getPointerInfo().getAddrSpace());
+  }
+
+  Register VecPtr = GLoadMI->getOperand(1).getReg();
+  LLT PtrTy = MRI.getType(VecPtr);
+
+  MachineFunction &MF = *MI.getMF();
+  auto *NewMMO = MF.getMachineMemOperand(&MMO, PtrInfo, VecEltVT);
+
+  LegalityQuery::MemDesc MMDesc(*NewMMO);
+
+  LegalityQuery Q = {TargetOpcode::G_LOAD, {VecEltVT, PtrTy}, {MMDesc}};
+
+  if (!isLegalOrBeforeLegalizer(Q))
+    return false;
+
+  // Load must be allowed and fast on the target.
+  LLVMContext &C = MF.getFunction().getContext();
+  auto &DL = MF.getDataLayout();
+  unsigned Fast = 0;
+  if (!getTargetLowering().allowsMemoryAccess(C, DL, VecEltVT, *NewMMO,
----------------
arsenm wrote:

You can still fold that into the same return statement 

https://github.com/llvm/llvm-project/pull/72670