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

[Pranav Taneja 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,
----------------
prtaneja wrote:

The isFast check is done similar to DAGCombiner equivalent of this function and this approach is also used in matchLoadOrCombine. 
Essentially, I think the allowsMemoryAccess() can modify the Fast variable depending on whether the access is fast or not.
Did you mean to take away the Fast variable and simplifying the check?

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