[AIEX] Align stack memsets to allow vector stores

In case of stack objects, if they are big enough, we
can improve alignment.

Author: andcarminati

llvm/lib/Target/AIE/AIECombine.td

@@ -222,6 +222,12 @@ def combine_phi_undef : GICombineRule<
          [{ return matchPhiOfUndef(*${root}, MRI, Observer, ${matchinfo}); }]),
   (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
 
+def combine_align_memset : GICombineRule<
+  (defs root:$root, build_fn_matchinfo:$matchinfo),
+  (match (wip_match_opcode G_MEMSET): $root,
+         [{ return matchAlignMemset(*${root}, MRI, (const AIEBaseInstrInfo &)B.getTII(), Observer, ${matchinfo}); }]),
+  (apply [{ Helper.applyBuildFnNoErase(*${root}, ${matchinfo}); }])>;
+
 // AIE-specifc combines (currently shared by AIE2 and AIE2P).
 def aie_additional_combines : GICombineGroup<[
   combine_unpad_vector,
@@ -240,7 +246,8 @@ def aie_additional_combines : GICombineGroup<[
   combine_concat_unmerge_phis,
   combine_trunc_ext,
   combine_load_const,
-  combine_phi_undef
+  combine_phi_undef,
+  combine_align_memset
 ]>;
 
 // AIE2P-specific combines.

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -20,6 +20,7 @@
 #include "llvm/CodeGen/GlobalISel/LegalizerHelper.h"
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegionInfo.h"
@@ -60,6 +61,10 @@ cl::opt<bool> CombineVecShiftByZero(
     "aie-combine-vec-shift-by-zero", cl::init(true), cl::Hidden,
     cl::desc("Combine vectors shift by zero into copies."));
 
+cl::opt<bool> MemsetOptimizations(
+    "aie-optimize-memsets", cl::init(true), cl::Hidden,
+    cl::desc("Apply memset optimizations (peeling/align/etc.)."));
+
 static unsigned getNumMaskUndefs(const ArrayRef<int> &Mask,
                                  unsigned StartIndex) {
   unsigned Count = 0;
@@ -3787,3 +3792,65 @@ bool llvm::matchPhiOfUndef(MachineInstr &Phi, MachineRegisterInfo &MRI,
 
   return true;
 }
+
+// If we have enough bytes set on a stack memset, we can simply
+// align this stack object to avoid store scalarization during
+// legalization.
+bool llvm::matchAlignMemset(MachineInstr &MI, MachineRegisterInfo &MRI,
+                            const AIEBaseInstrInfo &TII,
+                            GISelChangeObserver &Observer,
+                            BuildFnTy &MatchInfo) {
+  assert(MI.getOpcode() == TargetOpcode::G_MEMSET && "Expected a G_MEMSET");
+
+  if (!MemsetOptimizations)
+    return false;
+
+  // Try to keep alignment increase as minimum.
+  const unsigned BasicVectorByteSize = TII.getBasicVecRegSize() / 8;
+  // Half vectors are also supported.
+  const unsigned HalfVectorByteSize = BasicVectorByteSize / 2;
+
+  const Register CountReg = MI.getOperand(2).getReg();
+  auto Cst = getIConstantVRegValWithLookThrough(CountReg, MRI);
+  if (!Cst)
+    return false;
+  const unsigned ByteCount = Cst->Value.getZExtValue();
+
+  // Can we fill, at least, half of a basic vector?
+  if (ByteCount < HalfVectorByteSize)
+    return false;
+
+  const Register PtrReg = MI.getOperand(0).getReg();
+  const MachineInstr *DefDataInst = MRI.getUniqueVRegDef(PtrReg);
+
+  if (DefDataInst->getOpcode() != TargetOpcode::G_FRAME_INDEX)
+    return false;
+
+  if (MI.memoperands_empty())
+    return false;
+  MachineMemOperand *MMO = MI.memoperands().front();
+
+  const int FrameIndex = DefDataInst->getOperand(1).getIndex();
+
+  const Align OptimalAlign =
+      Align(ByteCount < BasicVectorByteSize ? HalfVectorByteSize
+                                            : BasicVectorByteSize);
+  const Align MMOAlign = MMO->getAlign();
+
+  if (MMOAlign == OptimalAlign)
+    return false;
+
+  MatchInfo = [=, &MI](MachineIRBuilder &B) {
+    MachineFunction *MF = MI.getMF();
+    MachineFrameInfo &MFI = MF->getFrameInfo();
+    MFI.setObjectAlignment(FrameIndex, OptimalAlign);
+    const LocationSize Size = MMO->getSize();
+    MI.dropMemRefs(*MF);
+    MI.addMemOperand(*MF,
+                     MF->getMachineMemOperand(
+                         MachinePointerInfo::getFixedStack(*MF, FrameIndex),
+                         MachineMemOperand::MOStore, Size, OptimalAlign));
+  };
+
+  return true;
+}

llvm/lib/Target/AIE/AIECombinerHelper.h

@@ -288,6 +288,10 @@ bool matchPhiBitcast(MachineInstr &Phi, MachineRegisterInfo &MRI,
 bool matchPhiOfUndef(MachineInstr &MI, MachineRegisterInfo &MRI,
                      GISelChangeObserver &Observer, BuildFnTy &MatchInfo);
 
+bool matchAlignMemset(MachineInstr &MI, MachineRegisterInfo &MRI,
+                      const AIEBaseInstrInfo &TII,
+                      GISelChangeObserver &Observer, BuildFnTy &MatchInfo);
+
 } // namespace llvm
 
 #endif