[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

@@ -4291,3 +4291,230 @@ 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.
+static bool areLoadUsesValidForExtractCombine(Register Reg,
+                                              unsigned ZExtExtractOpcode,
+                                              unsigned SExtExtractOpcode,
+                                              unsigned PadVectorOpcode,
+                                              MachineRegisterInfo &MRI) {
+
+  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 (UseOpcode == TargetOpcode::G_EXTRACT_VECTOR_ELT) {
+      // Direct extract is valid
+      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)) {
+        const unsigned PadUseOpcode = PadUse.getOpcode();
+        if (PadUseOpcode != ZExtExtractOpcode &&
+            PadUseOpcode != SExtExtractOpcode &&
+            PadUseOpcode != TargetOpcode::G_EXTRACT_VECTOR_ELT)
+          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 &MI, MachineRegisterInfo &MRI,
+                                     GISelChangeObserver &Observer,
+                                     BuildFnTy &MatchInfo) {
+  const MachineFunction &MF = *MI.getMF();
+  const AIEBaseInstrInfo &TII =
+      *static_cast<const AIEBaseInstrInfo *>(MF.getSubtarget().getInstrInfo());
+
+  const unsigned Opcode = MI.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 = MI.getOperand(2).getReg();
+  auto IdxCst = getIConstantVRegValWithLookThrough(IdxReg, MRI);
+  if (!IdxCst)
+    return false;
+  const int64_t Index = IdxCst->Value.getSExtValue();
+
+  // Get the vector operand
+  const Register VecReg = MI.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;
+
+  const unsigned ElemSizeInBytes = ElemSize / 8;
+
+  // 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 = [=, &MI, &MRI, &Observer](MachineIRBuilder &B) {
+    const Register PtrReg = LoadMI->getOperand(1).getReg();
+    const LLT S20 = LLT::scalar(20);
+
+    // Calculate byte offset: Index * ElemSizeInBytes
+    const int64_t ByteOffset = Index * ElemSizeInBytes;
+
+    // Set insertion point right after the original vector load
+    if (LoadMI->getNextNode())
+      B.setInstr(*LoadMI->getNextNode());
+    else
+      B.setInsertPt(*LoadMI->getParent(), LoadMI->getParent()->end());
+    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 If the byte offset is aligned to the original alignment, use
+    // it; otherwise use 1
+    const unsigned OrigAlign = MMO->getAlign().value();
+    const unsigned ScalarAlign = (ByteOffset % OrigAlign == 0) ? OrigAlign : 1;
+
+    // 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(MI);
+
+    // Handle the result based on the original opcode
+    if (IsZExtExtract || IsSExtExtract) {
+      // Need to extend to s32
+      const Register DstReg = MI.getOperand(0).getReg();
+      if (IsZExtExtract)
+        B.buildZExt(DstReg, LoadResultReg);
+      else
+        B.buildSExt(DstReg, LoadResultReg);
+    } else {
+      // G_EXTRACT_VECTOR_ELT - check if there's a G_ZEXT or G_SEXT user
+      const Register DstReg = MI.getOperand(0).getReg();
+      if (MRI.hasOneNonDBGUse(DstReg)) {
+        MachineInstr *UserMI = &*MRI.use_instr_nodbg_begin(DstReg);
+        const unsigned UserOpcode = UserMI->getOpcode();
+        if (UserOpcode == TargetOpcode::G_ZEXT ||
+            UserOpcode == TargetOpcode::G_SEXT) {
+          // Combine the extract and ext
+          const Register ExtDstReg = UserMI->getOperand(0).getReg();
+          if (UserOpcode == TargetOpcode::G_ZEXT)
+            B.buildZExt(ExtDstReg, LoadResultReg);
+          else
+            B.buildSExt(ExtDstReg, LoadResultReg);
+          Observer.erasingInstr(*UserMI);
+          UserMI->eraseFromParent();
+          Observer.erasingInstr(MI);
+          MI.eraseFromParent();
+          return;
+        }
+      }
+      // Just copy the result
+      B.buildCopy(DstReg, LoadResultReg);
+    }
+
+    Observer.erasingInstr(MI);
+    MI.eraseFromParent();
+  };
+
+  return true;
+}

llvm/lib/Target/AIE/AIECombinerHelper.h

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