[AIEX] Add a combiner to handle extract from unaligned vector load

In this case, we can load the scalar value directly instead of building
a full vector (legalizer will scalarize this load anyway) to extract.

Author: andcarminati

llvm/lib/Target/AIE/AIECombine.td

@@ -252,6 +252,12 @@ def combine_trunc_load : GICombineRule<
          [{ return matchNarrowTruncLoad(*${root}, MRI, Helper, Observer, ${matchinfo}); }]),
   (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
 
+def combine_unaligned_extract_load : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_EXTRACT_VECTOR_ELT, G_AIE_ZEXT_EXTRACT_VECTOR_ELT, G_AIE_SEXT_EXTRACT_VECTOR_ELT): $root,
+         [{ return matchUnalignedExtractLoad(*${root}, MRI, Observer, ${matchinfo}); }]),
+  (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
+
 // AIE-specifc combines (currently shared by AIE2 and AIE2P).
 def aie_additional_combines : GICombineGroup<[
   combine_unpad_vector,
@@ -274,7 +280,8 @@ def aie_additional_combines : GICombineGroup<[
   combine_align_memset,
   combine_peel_memset,
   combine_pack_stores_into_memset,
-  combine_trunc_load
+  combine_trunc_load,
+  combine_unaligned_extract_load
 ]>;
 
 // AIE2P-specific combines.
@@ -408,4 +415,3 @@ def AIE2PPostLegalizerCustomCombiner
                                                             combine_add_vector_elt_undef,
                                                              ]> {
 }
-

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -4292,3 +4292,210 @@ bool llvm::matchExtractVecEltAssertBcst(MachineInstr &MI,
 
   return true;
 }
+
+/// Helper function to recursively check if all uses of a register are valid
+/// for the unaligned extract load combiner.
+/// Automatically traverses through bitcasts to validate all usage patterns.
+/// Valid terminal uses are: direct extracts or pad vector operations (with use
+/// check).
+static bool areLoadUsesValidForExtractCombine(Register Reg,
+                                              unsigned ZExtExtractOpcode,
+                                              unsigned SExtExtractOpcode,
+                                              unsigned PadVectorOpcode,
+                                              MachineRegisterInfo &MRI) {
+
+  auto IsValidExtractOpcode = [&](unsigned Opcode) {
+    return Opcode == TargetOpcode::G_EXTRACT_VECTOR_ELT ||
+           Opcode == ZExtExtractOpcode || Opcode == SExtExtractOpcode;
+  };
+
+  for (const MachineInstr &Use : MRI.use_nodbg_instructions(Reg)) {
+    const unsigned UseOpcode = Use.getOpcode();
+
+    if (UseOpcode == TargetOpcode::G_BITCAST) {
+      // Recursively check bitcast uses
+      const Register BitcastDst = Use.getOperand(0).getReg();
+      if (!areLoadUsesValidForExtractCombine(BitcastDst, ZExtExtractOpcode,
+                                             SExtExtractOpcode, PadVectorOpcode,
+                                             MRI))
+        return false;
+      continue;
+    }
+
+    if (IsValidExtractOpcode(UseOpcode)) {
+      // Direct extract is valid (plain, zext, or sext)
+      continue;
+    }
+
+    if (UseOpcode == PadVectorOpcode) {
+      // Pad is valid if only used by extracts
+      const Register PadDst = Use.getOperand(0).getReg();
+      for (const MachineInstr &PadUse : MRI.use_nodbg_instructions(PadDst)) {
+        if (!IsValidExtractOpcode(PadUse.getOpcode()))
+          return false;
+      }
+      continue;
+    }
+
+    // Invalid use
+    return false;
+  }
+
+  return true;
+}
+
+/// Match unaligned vector loads that are only used for extracting elements
+/// and convert them to direct scalar loads.
+/// Supports s8, s16 and s32 element extractions from various vector
+/// configurations. Pattern:
+///   %vec:_(<N x sM>) = G_LOAD %ptr(p0) :: (align < M/8)
+///   %bitcast:_(<K x sX>) = G_BITCAST %vec
+///   %idx:_(s32) = G_CONSTANT i32 N
+///   %elt:_(sX) = G_EXTRACT_VECTOR_ELT %bitcast, %idx
+/// Or with G_AIE_PAD_VECTOR_UNDEF:
+///   %vec = G_LOAD %ptr :: (unaligned)
+///   %bitcast = G_BITCAST %vec
+///   %padded = G_AIE_PAD_VECTOR_UNDEF %bitcast
+///   %result:_(s32) = G_AIE_[Z/S]EXT_EXTRACT_VECTOR_ELT %padded, %idx
+/// Converts to:
+///   %offset:_(s20) = G_CONSTANT i20 (N * sizeof(sX))
+///   %new_ptr:_(p0) = G_PTR_ADD %ptr, %offset
+///   %elt:_(sX) = G_LOAD %new_ptr :: (align 1)
+///   %result:_(s32) = G_[Z/S]EXT %elt
+bool llvm::matchUnalignedExtractLoad(MachineInstr &ExtractMI,
+                                     MachineRegisterInfo &MRI,
+                                     GISelChangeObserver &Observer,
+                                     BuildFnTy &MatchInfo) {
+  const MachineFunction &MF = *ExtractMI.getMF();
+  const AIEBaseInstrInfo &TII =
+      *static_cast<const AIEBaseInstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  const unsigned Opcode = ExtractMI.getOpcode();
+  const unsigned ZExtExtractOpcode =
+      TII.getGenericExtractVectorEltOpcode(false);
+  const unsigned SExtExtractOpcode = TII.getGenericExtractVectorEltOpcode(true);
+  const unsigned PadVectorOpcode = TII.getGenericPadVectorOpcode();
+
+  const bool IsZExtExtract = (Opcode == ZExtExtractOpcode);
+  const bool IsSExtExtract = (Opcode == SExtExtractOpcode);
+  const bool IsPlainExtract = (Opcode == TargetOpcode::G_EXTRACT_VECTOR_ELT);
+
+  if (!IsZExtExtract && !IsSExtExtract && !IsPlainExtract)
+    return false;
+
+  // Get the index operand
+  const Register IdxReg = ExtractMI.getOperand(2).getReg();
+  const auto IdxCst = getIConstantVRegValWithLookThrough(IdxReg, MRI);
+  if (!IdxCst)
+    return false;
+  const int64_t Index = IdxCst->Value.getSExtValue();
+
+  // Get the vector operand
+  const Register VecReg = ExtractMI.getOperand(1).getReg();
+  const LLT VecTy = MRI.getType(VecReg);
+
+  // Check if vector has extractable element types (s8, s16, or s32)
+  if (!VecTy.isVector())
+    return false;
+
+  const LLT ElemTy = VecTy.getElementType();
+  const unsigned ElemSize = ElemTy.getSizeInBits();
+  if (ElemSize != 8 && ElemSize != 16 && ElemSize != 32)
+    return false;
+
+  // Trace back through G_AIE_PAD_VECTOR_UNDEF if present
+  MachineInstr *VecDefMI = MRI.getVRegDef(VecReg);
+  Register SourceVecReg = VecReg;
+
+  if (VecDefMI->getOpcode() == PadVectorOpcode) {
+    SourceVecReg = VecDefMI->getOperand(1).getReg();
+    VecDefMI = MRI.getVRegDef(SourceVecReg);
+  }
+
+  // Check for G_BITCAST (or direct vector if no bitcast needed)
+  Register LoadVecReg = SourceVecReg;
+  if (VecDefMI->getOpcode() == TargetOpcode::G_BITCAST)
+    LoadVecReg = VecDefMI->getOperand(1).getReg();
+
+  MachineInstr *LoadMI = MRI.getVRegDef(LoadVecReg);
+
+  // Check if it's a load
+  if (LoadMI->getOpcode() != TargetOpcode::G_LOAD)
+    return false;
+
+  // Check if the load is unaligned relative to the vector's total size
+  if (LoadMI->memoperands_empty())
+    return false;
+
+  const MachineMemOperand *MMO = LoadMI->memoperands().front();
+  const LLT LoadVecTy = MRI.getType(LoadVecReg);
+  const unsigned LoadVecSizeInBytes = LoadVecTy.getSizeInBytes();
+  // Vector is unaligned if alignment < vector size
+  // This allows extracting elements when the vector load itself is unaligned
+  if (MMO->getAlign().value() >= LoadVecSizeInBytes)
+    return false;
+
+  // Check that the loaded vector is only used by extracts (through bitcast and
+  // pad). The helper function will automatically traverse through bitcasts.
+  const Register LoadDstReg = LoadMI->getOperand(0).getReg();
+
+  if (!areLoadUsesValidForExtractCombine(LoadDstReg, ZExtExtractOpcode,
+                                         SExtExtractOpcode, PadVectorOpcode,
+                                         MRI))
+    return false;
+
+  // All checks passed, we can combine
+  MatchInfo = [=, &ExtractMI, &MRI, &Observer](MachineIRBuilder &B) {
+    const Register PtrReg = LoadMI->getOperand(1).getReg();
+    const LLT S20 = LLT::scalar(20);
+
+    const unsigned ElemSizeInBytes = ElemSize / 8;
+    const int64_t ByteOffset = Index * ElemSizeInBytes;
+
+    // Set insertion point right after the original vector load
+    B.setInsertPt(*LoadMI->getParent(), std::next(LoadMI->getIterator()));
+    B.setDebugLoc(LoadMI->getDebugLoc());
+
+    // Create offset constant and pointer add
+    const Register OffsetReg = B.buildConstant(S20, ByteOffset).getReg(0);
+    const Register NewPtrReg =
+        B.buildPtrAdd(MRI.getType(PtrReg), PtrReg, OffsetReg).getReg(0);
+
+    // Calculate alignment for scalar load based on original vector load
+    // alignment using GCD to find the maximum provable alignment
+    const unsigned OrigAlign = MMO->getAlign().value();
+    const unsigned ScalarAlign = std::gcd(OrigAlign, OrigAlign + ByteOffset);
+
+    // Create new scalar load with derived alignment
+    MachineFunction &MF = B.getMF();
+    MachineMemOperand *NewMMO =
+        MF.getMachineMemOperand(MMO->getPointerInfo(), MMO->getFlags(),
+                                ElemSizeInBytes, Align(ScalarAlign));
+
+    const Register LoadResultReg = MRI.createGenericVirtualRegister(ElemTy);
+    B.buildLoad(LoadResultReg, NewPtrReg, *NewMMO);
+
+    // Now set insertion point at the extract position for the copy/extension
+    B.setInstr(ExtractMI);
+
+    // Handle the result based on the original opcode
+    if (IsZExtExtract || IsSExtExtract) {
+      // Need to extend to s32
+      const Register DstReg = ExtractMI.getOperand(0).getReg();
+      if (IsZExtExtract)
+        B.buildZExt(DstReg, LoadResultReg);
+      else
+        B.buildSExt(DstReg, LoadResultReg);
+    } else {
+      // G_EXTRACT_VECTOR_ELT
+      const Register DstReg = ExtractMI.getOperand(0).getReg();
+      // Just copy the result
+      B.buildCopy(DstReg, LoadResultReg);
+    }
+
+    Observer.erasingInstr(ExtractMI);
+    ExtractMI.eraseFromParent();
+  };
+
+  return true;
+}

llvm/lib/Target/AIE/AIECombinerHelper.h

@@ -308,6 +308,11 @@ bool matchExtractVecEltAssertBcst(MachineInstr &MI, MachineRegisterInfo &MRI,
                                   const AIEBaseInstrInfo &TII,
                                   GISelChangeObserver &Observer,
                                   BuildFnTy &MatchInfo);
+
+bool matchUnalignedExtractLoad(MachineInstr &ExtractMI,
+                               MachineRegisterInfo &MRI,
+                               GISelChangeObserver &Observer,
+                               BuildFnTy &MatchInfo);
 } // namespace llvm
 
 #endif