Support for allowing direct VEXTRACT to 20-bit registers

llvm/lib/Target/AIE/AIE2InstrInfo.cpp

@@ -152,8 +152,9 @@ bool AIE2InstrInfo::verifyGenericInstruction(const MachineInstr &MI,
   switch (MI.getOpcode()) {
   case AIE2::G_AIE_ZEXT_EXTRACT_VECTOR_ELT:
   case AIE2::G_AIE_SEXT_EXTRACT_VECTOR_ELT:
-    ErrInfo = "Expected 32bit scalar destination";
-    return MRI.getType(MI.getOperand(0).getReg()) == LLT::scalar(32);
+    ErrInfo = "Expected 32bit or 20bit scalar destination";
+    return (MRI.getType(MI.getOperand(0).getReg()) == LLT::scalar(32) ||
+            MRI.getType(MI.getOperand(0).getReg()) == LLT::scalar(20));
   case AIE2::G_AIE_PAD_VECTOR_UNDEF:
     return verifySameLaneTypes(MI, ErrInfo) &&
            isLegalTypeToUnpad(MRI.getType(MI.getOperand(0).getReg()),
@@ -1417,6 +1418,22 @@ AIE2InstrInfo::getVExtractOpInfo(const MachineInstr &MI) const {
   }
 }
 
+bool AIE2InstrInfo::isS20Intrinsic(const MachineInstr &MI,
+                                   const MachineRegisterInfo &MRI) const {
+  const unsigned IntrinsicID = cast<GIntrinsic>(MI).getIntrinsicID();
+  switch (IntrinsicID) {
+  case Intrinsic::aie2_vextract_elem8_I512:
+  case Intrinsic::aie2_vextract_elem16_I512:
+  case Intrinsic::aie2_vextract_elem32_I512: {
+    // Check if the sign value is constant
+    const unsigned SignValueReg = MI.getOperand(4).getReg();
+    return getIConstantVRegSExtVal(SignValueReg, MRI).has_value();
+  }
+  default:
+    return false;
+  }
+}
+
 unsigned AIE2InstrInfo::getMaxLoadStoreSize() const { return 256; }
 
 bool AIE2InstrInfo::canCombineWithLoadStore(const MachineInstr &MI) const {

llvm/lib/Target/AIE/AIE2InstrInfo.h

@@ -180,6 +180,9 @@ class AIE2InstrInfo : public AIE2GenInstrInfo {
   std::optional<const VExtractOpInfo>
   getVExtractOpInfo(const MachineInstr &MI) const override;
 
+  bool isS20Intrinsic(const MachineInstr &MI,
+                      const MachineRegisterInfo &MRI) const override;
+
   unsigned getMaxLoadStoreSize() const override;
 
   bool canCombineWithLoadStore(const MachineInstr &MI) const override;

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -428,6 +428,12 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
     llvm_unreachable("Target didn't implement getVExtractOpInfo!");
   }
 
+  /// Check if the Intrinsic can produce S20
+  virtual bool isS20Intrinsic(const MachineInstr &MI,
+                              const MachineRegisterInfo &MRI) const {
+    llvm_unreachable("Target didn't implement isS20Intrinsic!");
+  }
+
   /// Return the maximun size for memory operations on this target.
   virtual unsigned getMaxLoadStoreSize() const {
     llvm_unreachable("Target didn't implement getMaxLoadStoreSize!");

llvm/lib/Target/AIE/AIECombinerHelper.cpp

@@ -581,6 +581,11 @@ static bool canProduceS20(const MachineRegisterInfo &MRI,
   case TargetOpcode::G_CONSTANT:
   case TargetOpcode::G_IMPLICIT_DEF:
     return true;
+  case TargetOpcode::G_INTRINSIC: {
+    const AIEBaseSubtarget &STI = AIEBaseSubtarget::get(*MI.getMF());
+    const AIEBaseInstrInfo *TII = STI.getInstrInfo();
+    return TII->isS20Intrinsic(MI, MRI);
+  }
   default:
     return false;
   }
@@ -842,6 +847,8 @@ bool modifyToS20(InstrNode Start, MachineRegisterInfo &MRI, MachineIRBuilder &B,
                  GISelChangeObserver &Observer, CombinerHelper &Helper) {
   const LLT S20 = LLT::scalar(20);
   MachineInstr *StartNodeMI = Start.getBaseNode();
+  const AIEBaseSubtarget &STI = AIEBaseSubtarget::get(*StartNodeMI->getMF());
+  const AIEBaseInstrInfo *TII = STI.getInstrInfo();
 
   // If Start can be rematerialized, only modify one user to use the
   // rematerialized instruction and leave the others unchanged.
@@ -901,16 +908,43 @@ bool modifyToS20(InstrNode Start, MachineRegisterInfo &MRI, MachineIRBuilder &B,
     Helper.tryCombineCopy(*StartNodeMI);
     return true;
   }
+  case TargetOpcode::G_INTRINSIC: {
+    if (TII->isS20Intrinsic(*StartNodeMI, MRI)) {
+      assert(StartNodeMI->getNumOperands() == 5);
+      Register ExtractDstReg = StartNodeMI->getOperand(0).getReg();
+      // Note: Operand 1 is the ID of the intrinsic
+      const Register SrcReg0 = StartNodeMI->getOperand(2).getReg();
+      const Register SrcReg1 = StartNodeMI->getOperand(3).getReg();
+      const Register SignReg = StartNodeMI->getOperand(4).getReg();
+
+      const auto SignVal = getIConstantVRegSExtVal(SignReg, MRI);
+      assert(SignVal.has_value() && "Expected SignVal to be constant");
+
+      B.setInsertPt(*StartNodeMI->getParent(), StartNodeMI);
+      Observer.changingInstr(*StartNodeMI);
+      MRI.setType(ExtractDstReg, S20);
+      auto NewMI =
+          B.buildInstr(TII->getGenericExtractVectorEltOpcode(SignVal.value()),
+                       {ExtractDstReg}, {SrcReg0, SrcReg1})
+              .getInstr();
+      Observer.changedInstr(*NewMI);
+
+      Helper.eraseInst(*StartNodeMI);
+      StartNodeMI = NewMI;
+    } else {
+      LLVM_DEBUG(dbgs() << "Node :" << *StartNodeMI);
+      llvm_unreachable("Unexpected G_INTRINSIC, while modifying IR");
+    }
+    break;
+  }
   default: {
     LLVM_DEBUG(dbgs() << "Node :" << *StartNodeMI);
     llvm_unreachable("Unexpected OpCode, while modifying IR");
   }
   }
 
-  switch (StartNodeMI->getOpcode()) {
-  case TargetOpcode::COPY:
-  case TargetOpcode::G_LOAD:
-  case TargetOpcode::G_PHI: {
+  // Function to handle the modification of instructions
+  auto modifyInstructionUses = [&](MachineInstr *StartNodeMI) {
     const auto UseInstIter =
         MRI.use_nodbg_instructions(StartNodeMI->getOperand(0).getReg());
     std::vector<MachineInstr *> UseInstr;
@@ -925,11 +959,27 @@ bool modifyToS20(InstrNode Start, MachineRegisterInfo &MRI, MachineIRBuilder &B,
       if (!modifyToS20(NextNodeToModify, MRI, B, Observer, Helper))
         llvm_unreachable("All input nodes should have updated");
     }
+  };
+
+  switch (StartNodeMI->getOpcode()) {
+  case TargetOpcode::COPY:
+  case TargetOpcode::G_LOAD:
+  case TargetOpcode::G_PHI: {
+    modifyInstructionUses(StartNodeMI);
     break;
   }
   default: {
-    LLVM_DEBUG(dbgs() << "Node :" << *StartNodeMI);
-    llvm_unreachable("Unexpected OpCode, while modifying IR");
+    const unsigned Opcode = StartNodeMI->getOpcode();
+    // Check if the opcode corresponds to a generic extract vector element
+    // operation. Since these opcodes are determined at runtime, we use an if
+    // statement to handle them.
+    if (Opcode == TII->getGenericExtractVectorEltOpcode(false) ||
+        Opcode == TII->getGenericExtractVectorEltOpcode(true)) {
+      modifyInstructionUses(StartNodeMI);
+    } else {
+      LLVM_DEBUG(dbgs() << "Node :" << *StartNodeMI);
+      llvm_unreachable("Unexpected OpCode, while modifying IR");
+    }
   }
   }
   return true;