[AIE2P] Transform VST.FLUSH into VST.FLUSH.CONV in PostSelectOptimize

Author: abhinay-anubola

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -259,6 +259,15 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
                            TypeSize Size) const {
     llvm_unreachable("Target didn't implement getCombinedPostIncOpcode");
   }
+
+  /// Check whether Opcode is a VST.PUSH.CONV
+  virtual bool isFifoStoreConvOpcode(unsigned Opcode) const { return false; }
+  /// \return Corresponding VST.FLUSH.CONV Opcode based on \a VST.FLUSH Opcode
+  virtual std::optional<unsigned>
+  getStoreFlushConvOpcode(unsigned StoreFlushOpcode) const {
+    llvm_unreachable("Target didn't implement getStoreFlushConvOpcode!");
+  }
+
   /// \return AIE2p OpCode based on \a IntrinsicID
   virtual unsigned getOpCode(MachineInstr &MI) const {
     llvm_unreachable("Target didn't implement getOpCode");

llvm/lib/Target/AIE/AIEPostSelectOptimize.cpp

@@ -584,6 +584,39 @@ bool fixLoadMemOpInfo(MachineFunction &MF, MachineBasicBlock &MBB,
   return Changed;
 }
 
+// Replace the VST.FLUSH opcode with VST.FLUSH.CONV if it is chained with
+// VST.PUSH.CONV. The CONV variant behaves identically to the normal variant but
+// all actions are delayed by one pipeline stage.
+bool modifyStoreFlush(MachineBasicBlock &MBB, MachineRegisterInfo &MRI) {
+  const TargetInstrInfo *TII = MBB.getParent()->getSubtarget().getInstrInfo();
+  const AIEBaseInstrInfo *AIEII = static_cast<const AIEBaseInstrInfo *>(TII);
+  bool Changed = false;
+
+  // Helper function to recursively update VST.FLUSH to VST.FLUSH.CONV
+  std::function<void(const Register)> Impl = [&](const Register UseReg) {
+    for (MachineInstr &UseMI : MRI.use_instructions(UseReg)) {
+      std::optional<unsigned> StoreFlushConvOpcode =
+          AIEII->getStoreFlushConvOpcode(UseMI.getOpcode());
+      if (StoreFlushConvOpcode) {
+        UseMI.setDesc(TII->get(*StoreFlushConvOpcode));
+        Changed = true;
+        // Update the opcode for the next dependent instruction in the chain
+        const Register UseDstReg = UseMI.getOperand(0).getReg();
+        Impl(UseDstReg);
+      }
+    }
+  };
+
+  for (MachineInstr &MI : MBB) {
+    if (AIEII->isFifoStoreConvOpcode(MI.getOpcode())) {
+      const Register DstReg = MI.getOperand(0).getReg();
+      Impl(DstReg);
+    }
+  }
+
+  return Changed;
+}
+
 bool AIEPostSelectOptimize::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "\n******* POST I-SEL OPTIMIZATION PASS *******\n"
                     << "********** Function: " << MF.getName() << '\n');
@@ -624,6 +657,14 @@ bool AIEPostSelectOptimize::runOnMachineFunction(MachineFunction &MF) {
     Changed |= fixLoadMemOpInfo(MF, MBB, MF.getRegInfo());
   }
 
+  // 5. Convert store flush instructions only on AIE2P targets: when VST.FLUSH
+  // and VST.PUSH.CONV are chained, replace VST.FLUSH with VST.FLUSH.CONV
+  if (MF.getTarget().getTargetTriple().isAIE2P()) {
+    for (MachineBasicBlock &MBB : MF) {
+      Changed |= modifyStoreFlush(MBB, MF.getRegInfo());
+    }
+  }
+
   return Changed;
 }
 

llvm/lib/Target/AIE/aie2p/AIE2PInstrInfo.cpp

@@ -270,6 +270,27 @@ bool AIE2PInstrInfo::isGenericOffsetMemOpcode(unsigned Opcode) const {
           (Opcode == AIE2P::G_AIE_OFFSET_ZEXTLOAD));
 }
 
+bool AIE2PInstrInfo::isFifoStoreConvOpcode(unsigned Opcode) const {
+  return ((Opcode == AIE2P::VST_PUSH_544_CONV_bfp16ebs16_ebs8) ||
+          (Opcode == AIE2P::VST_PUSH_544_CONV_bfp16ebs16_fp32) ||
+          (Opcode == AIE2P::VST_PUSH_576_CONV_bfp16ebs8_fp32));
+}
+
+std::optional<unsigned>
+AIE2PInstrInfo::getStoreFlushConvOpcode(unsigned StoreFlushOpcode) const {
+  switch (StoreFlushOpcode) {
+  case AIE2P::VST_FLUSH_512_normal_flush:
+    return AIE2P::VST_FLUSH_512_CONV_normal_flush;
+  case AIE2P::VST_FLUSH_512_fifo_1d_flush:
+    return AIE2P::VST_FLUSH_512_CONV_fifo_1d_flush;
+  case AIE2P::VST_FLUSH_512_2D:
+    return AIE2P::VST_FLUSH_512_CONV_2D;
+  case AIE2P::VST_FLUSH_512_3D:
+    return AIE2P::VST_FLUSH_512_CONV_3D;
+  }
+  return std::nullopt;
+}
+
 std::optional<unsigned> AIE2PInstrInfo::getCombinedPostIncOpcode(
     MachineInstr &BaseMemI, MachineInstr &PostIncI, TypeSize Size) const {
   switch (PostIncI.getOpcode()) {

llvm/lib/Target/AIE/aie2p/AIE2PInstrInfo.h

@@ -87,6 +87,7 @@ class AIE2PInstrInfo : public AIE2PGenInstrInfo {
   bool isBooleanNot(unsigned Opc) const override;
   bool isConstStep(const MachineInstr &MI, int64_t &Step) const override;
   bool isGenericOffsetMemOpcode(unsigned Opcode) const override;
+  bool isFifoStoreConvOpcode(unsigned Opcode) const override;
 
   bool verifyGenericInstruction(const MachineInstr &MI,
                                 StringRef &ErrInfo) const override;
@@ -97,6 +98,8 @@ class AIE2PInstrInfo : public AIE2PGenInstrInfo {
   std::optional<unsigned>
   getCombinedPostIncOpcode(MachineInstr &BaseMemI, MachineInstr &PtrAddI,
                            TypeSize Size) const override;
+  std::optional<unsigned>
+  getStoreFlushConvOpcode(unsigned StoreFlushOpcode) const override;
   unsigned getOpCode(MachineInstr &MI) const override;
   Register getVaddSignControlRegister() const override;