[AIE] Add inner-loop versioning to pipeline small-trip-count loops

llvm/include/llvm/IR/IntrinsicsAIE2.td

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2025 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 //
@@ -967,4 +967,12 @@ def int_aie2_extract_I128_I512 : ClangBuiltin<"__builtin_aiev2_extract_I128_I512
 // Set Sub-Vector
 def int_aie2_set_I512_I128 : ClangBuiltin<"__builtin_aiev2_set_I512_I128">, AIEV2Set_I512_I128;
 
+// Loop-versioning trip-count guard threshold used to pick the pipelined or
+// original copy of a versioned loop. The immediate is a placeholder, patched
+// by the post-pipeliner once the required minimum trip count is known.
+def int_aie2_loop_version_threshold
+    : DefaultAttrsIntrinsic<[llvm_i32_ty], [llvm_i32_ty],
+                            [IntrNoMem, IntrHasSideEffects, IntrNoDuplicate,
+                             ImmArg<ArgIndex<0>>]>;
+
 } // TargetPrefix = "aie2"

llvm/include/llvm/IR/IntrinsicsAIE2P.td

@@ -748,4 +748,10 @@ def int_aie2p_fifo_ld_pop_3d_unaligned: AIE2PFIFO_LD_POPUnaligned3D;
 def int_aie2p_fifo_ld_pop_576_3d_bfp16: AIE2PFIFO_LD_POP_576_BFP163D;
 def int_aie2p_fifo_ld_pop_544_3d_bfp16: AIE2PFIFO_LD_POP_576_BFP163D;
 
+// Loop-versioning trip-count guard threshold. See IntrinsicsAIE2.td.
+def int_aie2p_loop_version_threshold
+    : DefaultAttrsIntrinsic<[llvm_i32_ty], [llvm_i32_ty],
+                            [IntrNoMem, IntrHasSideEffects, IntrNoDuplicate,
+                             ImmArg<ArgIndex<0>>]>;
+
 } // TargetPrefix = "aie2p"

llvm/include/llvm/IR/IntrinsicsAIE2PS.td

@@ -1043,4 +1043,10 @@ def int_aie2ps_fifo_st_flush_2d_conv: AIE2PFIFO_ST_FLUSH_2D;
 def int_aie2ps_fifo_st_flush_3d: AIE2PFIFO_ST_FLUSH_3D;
 def int_aie2ps_fifo_st_flush_3d_conv: AIE2PFIFO_ST_FLUSH_3D;
 
+// Loop-versioning trip-count guard threshold. See IntrinsicsAIE2.td.
+def int_aie2ps_loop_version_threshold
+    : DefaultAttrsIntrinsic<[llvm_i32_ty], [llvm_i32_ty],
+                            [IntrNoMem, IntrHasSideEffects, IntrNoDuplicate,
+                             ImmArg<ArgIndex<0>>]>;
+
 } // TargetPrefix = "aie2ps"

llvm/lib/Target/AIE/AIE.h

@@ -120,6 +120,13 @@ llvm::FunctionPass *createReservedRegsLICMPass();
 extern char &AIEOuterLoopPipelinerID;
 void initializeAIEOuterLoopPipelinerPass(PassRegistry &);
 llvm::FunctionPass *createAIEOuterLoopPipelinerPass();
+
+// Inner Loop Versioning (IR-level). Emits a runtime trip-count guard around a
+// pipelined copy of a single-block inner loop whose minimum trip count is too
+// small for the software pipeliner. See AIEInnerLoopVersioning.cpp.
+extern char &AIEInnerLoopVersioningID;
+void initializeAIEInnerLoopVersioningPass(PassRegistry &);
+llvm::FunctionPass *createAIEInnerLoopVersioningPass();
 } // namespace llvm
 
 #endif

llvm/lib/Target/AIE/AIE2InstrInfo.cpp

@@ -1415,6 +1415,10 @@ AIE2InstrInfo::getZOLSupport() const {
   return Result;
 }
 
+std::optional<unsigned> AIE2InstrInfo::getLoopVersionThresholdOpcode() const {
+  return AIE2::PseudoLoopVersionThreshold;
+}
+
 std::optional<AIEBaseInstrInfo::JNZDSupport>
 AIE2InstrInfo::getJNZDSupport() const {
   AIEBaseInstrInfo::JNZDSupport Result;

llvm/lib/Target/AIE/AIE2InstrInfo.h

@@ -106,6 +106,7 @@ class AIE2InstrInfo : public AIE2GenInstrInfo {
 
   virtual std::optional<ZOLSupport> getZOLSupport() const override;
   virtual std::optional<JNZDSupport> getJNZDSupport() const override;
+  std::optional<unsigned> getLoopVersionThresholdOpcode() const override;
   virtual std::optional<IfConvSupport> getIfConvSupport() const override;
 
   virtual bool

llvm/lib/Target/AIE/AIE2InstructionSelector.cpp

@@ -337,6 +337,8 @@ bool AIE2InstructionSelector::select(MachineInstr &I) {
       return selectStartLoop(I, MRI);
     case Intrinsic::set_loop_iterations:
       return selectSetLoopIterations(I, MRI, MIB);
+    case Intrinsic::aie2_loop_version_threshold:
+      return selectLoopVersionThreshold(I, MRI, MIB);
     case Intrinsic::aie2_scd_read_vec:
     case Intrinsic::aie2_scd_read_acc32:
     case Intrinsic::aie2_scd_expand_lo:

llvm/lib/Target/AIE/AIE2MultiSlotPseudoInstrInfo.td

@@ -4,7 +4,7 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
-// (c) Copyright 2023-2024 Advanced Micro Devices, Inc. or its affiliates
+// (c) Copyright 2023-2026 Advanced Micro Devices, Inc. or its affiliates
 //
 //===----------------------------------------------------------------------===//
 
@@ -71,6 +71,18 @@ let isMoveImm = 1, isReMaterializable = 1, isAsCheapAsAMove = 1, Itinerary = II_
       [MOVA_lda_cg, MOVXM_lng_cg]>;
 }
 
+// Loop-versioning trip-count guard threshold, produced by the AIE loop
+// versioning pass. The post-pipeliner patches its immediate operand with the
+// required minimum trip count once a modulo schedule is found.
+// Side-effecting and non-duplicable so MachineCSE never merges sibling guards'
+// placeholders; dest is eR (the guard compare's class) so MOVA can be used.
+let isMoveImm = 1, Itinerary = II_MOV, hasSideEffects = 1, isNotDuplicable = 1,
+    mayLoad = 0, mayStore = 0 in
+  def PseudoLoopVersionThreshold
+      : MultiSlot_Pseudo<(outs eR:$dst), (ins simm10:$i),
+                         "pseudo_loop_version_threshold", "$dst, $i",
+                         [MOVA_lda_cg, MOV_mv_cg, MOVXM_lng_cg]>;
+
 // Pseudo VLD
 let hasSideEffects = false, mayLoad = true, mayStore = false in {
   let Itinerary = II_VLDA_W in {

llvm/lib/Target/AIE/AIE2TargetMachine.cpp

@@ -79,6 +79,7 @@ TargetPassConfig *AIE2TargetMachine::createPassConfig(PassManagerBase &PM) {
 
 bool AIE2PassConfig::addPreISel() {
   if (TM->getOptLevel() != CodeGenOptLevel::None) {
+    addPass(createAIEInnerLoopVersioningPass());
     addPass(createHardwareLoopsLegacyPass());
     addPass(createAIEOuterLoopPipelinerPass());
   }

llvm/lib/Target/AIE/AIEBaseInstrInfo.cpp

@@ -39,6 +39,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
 #include <limits>
 
 #define DEBUG_TYPE "aie-codegen"
@@ -424,6 +425,26 @@ void AIEBaseInstrInfo::adjustTripCount(MachineInstr &MI, int Adjustment) const {
   Imm.setImm(Imm.getImm() + Adjustment);
 }
 
+bool AIEBaseInstrInfo::isLoopVersionThresholdDef(const MachineInstr &MI) const {
+  auto Opcode = getLoopVersionThresholdOpcode();
+  return Opcode && MI.getOpcode() == *Opcode;
+}
+
+void AIEBaseInstrInfo::setLoopVersionThreshold(MachineInstr &MI,
+                                               int MinTripCount) const {
+  assert(isLoopVersionThresholdDef(MI));
+  // The threshold materializes into a narrow scalar-move immediate; enforce the
+  // narrowest field width (AIE2's simm10) across subtargets. A stage count is
+  // tiny in practice, so this only trips on a pathological schedule. Checked
+  // unconditionally, not just under assert: a truncated threshold would select
+  // the high-trip-count copy for trip counts too small for the schedule, a
+  // miscompile.
+  if (!isInt<10>(MinTripCount))
+    report_fatal_error("AIE loop versioning: threshold does not fit the "
+                       "scalar-move immediate");
+  MI.getOperand(1).setImm(MinTripCount);
+}
+
 bool AIEBaseInstrInfo::isHardwareLoopStart(unsigned Opcode) const {
   const auto ZOLSupport = getZOLSupport();
   return ZOLSupport && Opcode == ZOLSupport->LoopStartOpcode;

llvm/lib/Target/AIE/AIEBaseInstrInfo.h

@@ -520,6 +520,22 @@ struct AIEBaseInstrInfo : public TargetInstrInfo {
   /// \pre isZOLTripCountDef(MI)
   virtual void adjustTripCount(MachineInstr &MI, int Update) const;
 
+  /// Opcode of this subtarget's loop-versioning threshold pseudo
+  /// (PseudoLoopVersionThreshold), or std::nullopt if unsupported. The pseudo
+  /// defines the threshold register (operand 0) and holds the patchable
+  /// threshold immediate (operand 1); see its def for the design rationale.
+  virtual std::optional<unsigned> getLoopVersionThresholdOpcode() const {
+    return std::nullopt;
+  }
+
+  /// Check whether \p MI is the loop-versioning threshold move.
+  bool isLoopVersionThresholdDef(const MachineInstr &MI) const;
+
+  /// Set the guard threshold immediate on the loop-versioning threshold
+  /// move \p MI to the (absolute) minimum trip count \p MinTripCount.
+  /// \pre isLoopVersionThresholdDef(MI)
+  void setLoopVersionThreshold(MachineInstr &MI, int MinTripCount) const;
+
   /// Check whether this is a zero-overhead loop start block
   virtual bool isZeroOverheadLoopSetupInstr(const MachineInstr &) const;
 

llvm/lib/Target/AIE/AIEBaseInstructionSelector.cpp

@@ -142,6 +142,18 @@ bool AIEBaseInstructionSelector::selectSetLoopIterations(
   return constrainSelectedInstRegOperands(*LS, TII, TRI, RBI);
 }
 
+bool AIEBaseInstructionSelector::selectLoopVersionThreshold(
+    MachineInstr &I, MachineRegisterInfo &MRI, MachineIRBuilder &MIB) {
+  auto Opcode = TII.getLoopVersionThresholdOpcode();
+  assert(Opcode && "Loop versioning not supported on this subtarget");
+
+  const Register DstReg = I.getOperand(0).getReg();
+  const int64_t Threshold = I.getOperand(2).getImm();
+  auto Mov = MIB.buildInstr(*Opcode, {DstReg}, {}).addImm(Threshold);
+  I.eraseFromParent();
+  return constrainSelectedInstRegOperands(*Mov, TII, TRI, RBI);
+}
+
 // Try to match BRCOND(Intrinsic::loop_decrement)
 bool AIEBaseInstructionSelector::selectBrCondLoopDecrement(
     MachineInstr &BrCond, MachineRegisterInfo &MRI) {

llvm/lib/Target/AIE/AIEBaseInstructionSelector.h

@@ -155,6 +155,8 @@ class AIEBaseInstructionSelector : public InstructionSelector {
                  MachineIRBuilder &MIB, MachineInstrBuilder &MI);
   bool selectSetLoopIterations(MachineInstr &I, MachineRegisterInfo &MRI,
                                MachineIRBuilder &MIB);
+  bool selectLoopVersionThreshold(MachineInstr &I, MachineRegisterInfo &MRI,
+                                  MachineIRBuilder &MIB);
   bool selectBrCondLoopDecrement(MachineInstr &BrCond,
                                  MachineRegisterInfo &MRI);
   bool selectBrCondLoopDecrementReg(MachineInstr &BrCond,

llvm/lib/Target/AIE/AIEBaseTargetMachine.cpp

@@ -177,6 +177,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeAIETarget() {
   initializeAIERegClassConstrainerPass(*PR);
   initializeReservedRegsLICMPass(*PR);
   initializeAIEOuterLoopPipelinerPass(*PR);
+  initializeAIEInnerLoopVersioningPass(*PR);
   initializeAIESpillSlotOptimizationPass(*PR);
 }