[mlgo] Add in-development instruction based features for regalloc advisor

[aidengrossman]

Summary

This patch adds in instruction based features to the regalloc advisor
gated behind a flag so a user can decide at runtime whether or not they
want to enable the feature. The features are only enabled when LLVM is
compiled in MLGO develpment mode (LLVM_HAVE_TF_API) is set to true.

To extract the instruction features, I'm taking a list of segments from
each LiveInterval and noting the start and end SlotIndices. This list is then
sorted based on the start SlotIndex and I iterate through each SlotIndex
to grab instructions, making sure to check for overlaps. This results in
a vector of opcodes and binary mapping matrix that maps live ranges to the
opcodes of the instructions within that LR.

[mtrofin]

llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll

3

Because the output very verbose, perhaps just checking specific values may be more maintainable? Like checking that you counted things correctly in your 33x300 matrix - add some comments maybe about what the output looks like and why certain sentinel values are expected. This would also help one identify, later, if the test breaks legitimately due to a change (e.g. some machine instruction is now not present anymore -> of course the features change) vs due to a bug.

[aidengrossman]

Changed up dev mode extra features test to make it less fragile by getting rid
of the exact logging diff check and adding more FileCheck checks on the
instructions and mapping matrix for the first eviction problem.

[aidengrossman]

llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll

3

That's a very good point. The test could potentially be pretty fragile if checking the exact code output. I reworked the check to get rid of the exact diff and added a lot more FileCheck checks along with comments to make sure that the output looks as expected to some reasonable degree.

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

925

this deserves more comments about: the whole structure of the data packet you're preparing and the various steps along the way - it'll greatly help with maintainability, etc.

926

may be more readable to have a small struct with the 2 slot indices and size_t with nice readable names, and then avoid the whole std::get<x> business.

942

nit: we generally do if (!CurrentMachineInstruction).

1033

nit: or, // continue from the feature index the previous loop left off ?

llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll

3

thanks, this looks way more manageable!

52

add a new line

[aidengrossman]

Adjust comments to better document the extraction algorithm and the data that
it returns. Also refactor to passing around a struct rather than a tuple and
adjust a comment in the dev mode logging.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

925

I've added quite a few more comments on the structure of the extraction algorithm itself and the structure of the data that it extracts. Let me know if anything is unclear or too verbose and I'll work on fixing it.

1033

This was left over from some experimentation that I have since removed and thus was commenting a line that isn't even there. I've changed it to note your suggestion and also to explain why the indexing is slightly odd (only going to FeaturesWithDevelopmentCount -1 1 rather than FeaturesWithDevelopmentCount itself.

[aidengrossman]

Fixed release mode build and gated more of the development feature tooling
with the preprocessing behind the LLVM_HAVE_TF_API flag.

[aidengrossman]

Just making sure this patch doesn't go anywhere yet. When built with assertions on, the new development features are currently tripping an assertion when CurrentIndex.getNextIndex() is called when CurrentMachineInstruction is equal to nullptr. Need to do some more debugging/investigation see why this is going on.

[aidengrossman]

Add if statement to handle edge case where the current machine instruction
is null and it tries to skip to the slotindex in the next machine instruction
but can't due to it being at the of the slotindex analysis.

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

149

rename this to OpcodeValueCutoff - Count to me refers at this point to how many opcodes in the tensor, but IIUC you actually mean values larger than that are unknown to us. Same for the comment.

265

you can avoid the -1 thing by splitting your RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST in 2, the "first" and the "rest", and then forcing the index of the first to be == FeatureCount

#ifdef LLVM_HAVE_TF_API
#define RA_FIRST_EXTRA_FEATURE(M)  \
    M(int64_t, instructions, InstructionsShape,                                  \
    "Opcodes of the instructions covered by the eviction problem") 
#define RA_REST_EXTRA_FEATURES(M)                            \
  M(int64_t, instructions_mapping, InstructionsMappingShape,                   \
    "A binary matrix mapping LRs to instruction opcodes")
#define RA_EVICT_FEATURES_UNDER_DEVELOPMENT(M) \
  RA_FIRST_EXTRA_FEATURE(M) \
  RA_REST_EXTRA_FEATURES(M)
#else
#define RA_FIRST_EXTRA_FEATURE(M)
#define   RA_REST_EXTRA_FEATURES(M)
#endif

#define RA_EVICT_FEATURES_UNDER_DEVELOPMENT(M) \
  RA_FIRST_EXTRA_FEATURE(M) \
  RA_REST_EXTRA_FEATURES(M)
enum FeatureIDs {
#define _FEATURE_IDX_SIMPLE(_, name, __, ___) name
#define _FEATURE_IDX(A,B,C,D) _FEATURE_IDX_SIMPLE(A,B,C,D),
  RA_EVICT_FEATURES_LIST(_FEATURE_IDX) FeatureCount,
  RA_FIRST_EXTRA_FEATURE(_FEATURE_IDX_SIMPLE) = FeatureCount,
  RA_REST_EXTRA_FEATURES(_FEATURE_IDX)
      FeaturesWithDevelopmentCount
#undef _FEATURE_IDX
};

286

Initialize things that don't have default ctors easy to do at def, and avoids nondeterministic bugs. Like Pos, at minimum, set it to 0.

288

add a comment as to what Pos is - it's the index of the column corresponding to the physical register of the live interval segment captured by this LRStartEndInfo, right? (or the candidate)

also a comment for LRStartEndInfo overall

999

This feels like it'd benefit from a unittest. The core logic is about intervals that cover instruction opcodes, so the only LLVM-ness of it is in LIS->getInstructionFromIndex(CurrentIndex) (because LIS->getSlotIndexes()->getLastIndex()) is basically a constant you can pass in)

You can test it by making a small change: make this into a utility that takes LRPosInfo and a std::function (or a llvm::function_ref, whatever) that gives the opcode for a SlotIndex. The utility doesn't need to know about mlevictadvisor or anything - it's just dealing with intervals. You can also pass a Runner in, and in the test case, just use the NoInferenceModelRunner.

So then you can set up all sorts of interesting interval overlapping cases, and you can just populate your opcodes with incrementing numbers or something.

The nice thing is that the unittest doesn't need any #define specific stuff - it's generic.

The rest can stay the same - meaning, let MLEvictAdvisor::extractInstructionFeatures call that utility, etc.

1004

(nit - it took me a few reads to get it) "A vector of size "max instruction count". It contains the instruction opcodes of instructions covered by all intervals in LRPosInfo"

wdyt?

1013

This is rather the Instruction "Index" or something like that, right? Count to me means total.

1014

Nit: CurrentSegmentIdx

1020

typo: process"ed"

1043

you don't need the -1 thing here if you split the extra list in 2

1058

the exit condition could include LRPosInfo[OverlapCheckCurrentSegment].Begin > CurrentIndex, for more clarity?

1068

can't LRPosInfo[OverlapCheckCurrentSegment].End < CurrentIndex? So we know the CurrentIndex is below Begin, but it could also be below this segment's end?

Current instruction: 5
CurrentSegment: [1,7)
NextSegment: [2,4)
NextNextSegment: [2, 8)

1084

you just need to go to the next segment and potentially set CurrentIndex to that one's beginning if it's after CurrentIndex (simpler code)

[aidengrossman]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

286

That's a good point. Done for Pos. The SlotIndex class has a default constructor creating an invalid index, so we should be good there.

288

Added a comment for the entirety of LRStartEndInfo which also covers the meaning of the Pos member.

[aidengrossman]

Address some of the inline comments (still a couple more to get to).

[aidengrossman]

Addressed more inline comments. Still some to go.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1068

LRPosInfo[OverlapCheckCurrentSegment].End can be less than CurrentIndex, but it shouldn't be missing any instructions. It's not checking if CurrentIndex is below begin, but rather if CurrentIndex is above the beginning of the current segment, because if the beginning of the current segment being checked is greater than the current index under analysis, it is guaranteed that the segment currently being processed as well as all future segments don't contain the currently being processed index due to the segments being sorted in ascending order by beginning index.

[aidengrossman]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

1068

Nevermind my previous comment. Not sure what I was thinking there. This was an actual issue and I ended up catching it while looking through some unit test cases. Should be fixed on the next push as I added in a condition to only set the relevant element of the instructions_mapping matrix if CurrentIndex <= LRPosInfo[OverlapCheckCurrentSegment].End.

[aidengrossman]

Refactor for unit tests, add unit tests, and modify main feature extraction code
to fix some edge cases found through unit testing.

[aidengrossman]

Removed conditional compilation for the feature extraction for the in
development features and moved size constants to MLRegallocEvictAdvisor.h to
allow for the constants to be used during unit testing.

[mtrofin]

llvm/include/llvm/CodeGen/SlotIndexes.h

148 ↗

(On Diff #460937)

(just noticed this) can you avoid this, because it's just for test, and instead just replace the SlotIndex in your collection?

[aidengrossman]

llvm/include/llvm/CodeGen/SlotIndexes.h

148 ↗

(On Diff #460937)

This is for instantiation of each SlotIndex within the list. I would like to avoid it if possible too, but in order to setup the test case the SlotIndex needs to have a valid IndexListEntry* pointer, and there is no publicly exposed constructor that allows setting that, so I'm calling the default constructor and setting the pointer with this function. So it doesn't seem like there's a super practical way to avoid this. Unless there's some easy way to avoid this that I'm just not seeing?

[mtrofin]

llvm/include/llvm/CodeGen/SlotIndexes.h

148 ↗

(On Diff #460937)

Ugh... I'd rather the ctor be public - that way, the state management of a SlotIndex stays the same (once created, it doesn't change its lie - easier to understand / maintain)

or (my initial preference) introduce an abstraction around the slot index stuff. IIRC there isn't that much you actually need from that API.

[aidengrossman]

Changed test setup to use SlotIndex constructor instead of setting the
ListIndexEntry pointer directly by getting rid of the pointer setter
and making the constructor public with a comment nothing that the
constructor isn't for general public consumption.

[aidengrossman]

Fixed conditional compilation for dev features so that the default case
doesn't break.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

 //===- MLRegAllocEvictAdvisor.cpp - ML eviction advisor -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 static cl::opt<std::string> TrainingLog(
     "regalloc-training-log", cl::Hidden,
     cl::desc("Training log for the register allocator eviction model"));
 
 static cl::opt<std::string> ModelUnderTraining(
     "regalloc-model", cl::Hidden,
     cl::desc("The model being trained for register allocation eviction"));
 
+static cl::opt<bool> EnableDevelopmentFeatures(
+    "regalloc-enable-development-features", cl::Hidden,
+    cl::desc("Whether or not to enable features under development for the ML "
+             "regalloc advisor"));
+
+#else
+static const bool EnableDevelopmentFeatures = false;
 #endif // #ifdef LLVM_HAVE_TF_API
 
 extern cl::opt<unsigned> EvictInterferenceCutoff;
 
 /// The score injection pass.
 /// This pass calculates the score for a function and inserts it in the log, but
 /// this happens only in development mode. It's a no-op otherwise.
 namespace llvm {

 // matrix. The rows are the features (see next). The columns correspond to the
 // interferences. We treat the candidate virt reg as an 'interference', too, as
 // its feature set is the same as that of the interferring ranges. So we'll have
 // MaxInterferences + 1 columns and by convention, we will use the last column
 // for the virt reg seeking allocation.
 static const int64_t CandidateVirtRegPos = MaxInterferences;
 static const int64_t NumberOfInterferences = CandidateVirtRegPos + 1;
 
+// The number of instructions that a specific live range might have is variable,
+// but we're passing in a single matrix of instructions and tensorflow saved
+// models only support a fixed input size, so we have to cap the number of
+// instructions that can be passed along. The specific value was derived from
+// experimentation such that the majority of eviction problems would be
+// completely covered.
+static const int ModelMaxSupportedInstructionCount = 300;
+static const std::vector<int64_t> InstructionsShape{
+    1, ModelMaxSupportedInstructionCount};
+static const std::vector<int64_t> InstructionsMappingShape{
+    1, NumberOfInterferences, ModelMaxSupportedInstructionCount};
+
+// The model can only accept a specified number of opcodes and will error it if
+// fed an opcode it hasn't seen before. This constant sets the current cutoff.
+static const int OpcodeCountCutoff = 17716;
+
 // Most features are as described above, so we'll reuse this vector in defining
 // them.
 static const std::vector<int64_t> PerLiveRangeShape{1, NumberOfInterferences};
 
 // --------------
 // Features table
 // --------------
 // For each interfering live range (incl. the candidate) we collect a number of
#define RA_EVICT_FEATURES_LIST(M)                                              \
   M(float, use_def_density, PerLiveRangeShape,                                 \
     "the max weight, as computed by the manual heuristic")                     \
   M(int64_t, max_stage, PerLiveRangeShape,                                     \
     "largest stage of an interval in this LR")                                 \
   M(int64_t, min_stage, PerLiveRangeShape,                                     \
     "lowest stage of an interval in this LR")                                  \
   M(float, progress, {1}, "ratio of current queue size to initial size")
 
+#ifdef LLVM_HAVE_TF_API
+#define RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(M)                            \
+  M(int64_t, instructions, InstructionsShape,                                  \
+    "Opcodes of the instructions covered by the eviction problem")             \
+  M(int64_t, instructions_mapping, InstructionsMappingShape,                   \
+    "A binary matrix mapping LRs to instruction opcodes")
+#else
+#define RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(M)
+#endif
+
 // The model learns to pick one of the mask == 1 interferences. This is the
 // name of the output tensor. The contract with the model is that the output
 // will be guaranteed to be to a mask == 1 position. Using a macro here to
 // avoid 'not used' warnings (and keep cond compilation to a minimum)
 #define DecisionName "index_to_evict"
 
 // Named features index.
 enum FeatureIDs {
 #define _FEATURE_IDX(_, name, __, ___) name,
-  RA_EVICT_FEATURES_LIST(_FEATURE_IDX)
+  RA_EVICT_FEATURES_LIST(_FEATURE_IDX) FeatureCount,
+  RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(_FEATURE_IDX)
+      FeaturesWithDevelopmentCount
 #undef _FEATURE_IDX
-      FeatureCount
 };
 
 // The ML advisor will typically have a sparse input to the evaluator, because
 // various phys regs won't be available. It's easier (maintenance-wise) to
 // bulk-reset the state of the evaluator each time we are about to use it
 // again.
 template <typename T> size_t getTotalSize(const std::vector<int64_t> &Shape) {
   size_t Ret = sizeof(T);
   for (const auto V : Shape)
     Ret *= V;
   return Ret;
 }
 
 void resetInputs(MLModelRunner &Runner) {
 #define _RESET(TYPE, NAME, SHAPE, __)                                          \
   std::memset(Runner.getTensorUntyped(FeatureIDs::NAME), 0,                    \
               getTotalSize<TYPE>(SHAPE));
   RA_EVICT_FEATURES_LIST(_RESET)
 #undef _RESET
+  // when resetting the features under development we need to make sure to skip
+  // FeatureIDs::FeatureCount
+  if (EnableDevelopmentFeatures) {
+#define _RESET_DEV(TYPE, NAME, SHAPE, __)                                      \
+  std::memset(Runner.getTensorUntyped(FeatureIDs::NAME - 1), 0,                \
+              getTotalSize<TYPE>(SHAPE));
+    RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(_RESET_DEV)
+  }
+#undef _RESET_DEV
 }
 
 // Per-live interval components that get aggregated into the feature values
 // that will be passed to the evaluator.
 struct LIFeatureComponents {
   double R = 0;
   double W = 0;
   double RW = 0;
   double IndVarUpdates = 0;
   double HintWeights = 0.0;
   int64_t NrDefsAndUses = 0;
   float HottestBlockFreq = 0.0;
   bool IsRemat = false;
 };
 
 using CandidateRegList =
     std::array<std::pair<MCRegister, bool>, NumberOfInterferences>;
 using FeaturesListNormalizer =
     llvm::SmallVector<float, FeatureIDs::FeatureCount>;
 
 /// The ML evictor (commonalities between release and development mode)
 class MLEvictAdvisor : public RegAllocEvictionAdvisor {
 public:
   MLEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
                  MLModelRunner *Runner, const MachineBlockFrequencyInfo &MBFI,
                  const MachineLoopInfo &Loops);
 
protected:
   virtual int64_t
   tryFindEvictionCandidatePosition(const LiveInterval &VirtReg,
                                    const AllocationOrder &Order,
                                    unsigned OrderLimit, uint8_t CostPerUseLimit,
                                    const SmallVirtRegSet &FixedRegisters) const;
 
   /// Load the features of the given VirtReg (allocated or not) at column Pos,
   /// but if  that can't be evicted, return false instead.
-  bool loadInterferenceFeatures(const LiveInterval &VirtReg, MCRegister PhysReg,
-                                bool IsHint,
-                                const SmallVirtRegSet &FixedRegisters,
-                                llvm::SmallVectorImpl<float> &Largest,
-                                size_t Pos) const;
+  bool loadInterferenceFeatures(
+      const LiveInterval &VirtReg, MCRegister PhysReg, bool IsHint,
+      llvm::SmallVectorImpl<float> &Largest, size_t Pos,
+      SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>> &LRPosInfo)
+      const;
 
 private:
   static float getInitialQueueSize(const MachineFunction &MF);
 
   MCRegister tryFindEvictionCandidate(
       const LiveInterval &VirtReg, const AllocationOrder &Order,
       uint8_t CostPerUseLimit,
       const SmallVirtRegSet &FixedRegisters) const override;
 
   void extractFeatures(const SmallVectorImpl<const LiveInterval *> &Intervals,
                        llvm::SmallVectorImpl<float> &Largest, size_t Pos,
-                       int64_t IsHint, int64_t LocalIntfsCount,
-                       float NrUrgent) const;
+                       int64_t IsHint, int64_t LocalIntfsCount, float NrUrgent,
+                       SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>>
+                           &LRPosInfo) const;
+
+  void extractInstructionFeatures(
+      llvm::SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>>
+          &LRPosInfo) const;
 
   // Point-in-time: we didn't learn this, so we always delegate to the
   // default.
   bool canEvictHintInterference(
       const LiveInterval &VirtReg, MCRegister PhysReg,
       const SmallVirtRegSet &FixedRegisters) const override {
     return getDefaultAdvisor().canEvictHintInterference(VirtReg, PhysReg,
                                                         FixedRegisters);

 // ===================================
 // Release (AOT) - specifics
 // ===================================
 class ReleaseModeEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {
 public:
   ReleaseModeEvictionAdvisorAnalysis()
       : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {
+    if (EnableDevelopmentFeatures) {
+      InputFeatures = {
+          RA_EVICT_FEATURES_LIST(_DECL_FEATURES)
+              RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(_DECL_FEATURES)};
+    } else {
       InputFeatures = {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)};
     }
+  }
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Release;
   }
 
 private:
   std::vector<TensorSpec> InputFeatures;
 
private:
   Logger *const Log;
 };
 
 class DevelopmentModeEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {
 public:
   DevelopmentModeEvictionAdvisorAnalysis()
       : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Development) {
+    if (EnableDevelopmentFeatures) {
+      InputFeatures = {
+          RA_EVICT_FEATURES_LIST(_DECL_FEATURES)
+              RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(_DECL_FEATURES)};
+      TrainingInputFeatures = {
+          RA_EVICT_FEATURES_LIST(_DECL_TRAIN_FEATURES)
+              RA_EVICT_FEATURES_UNDER_DEVELOPMENT_LIST(_DECL_FEATURES)
+                  TensorSpec::createSpec<float>("action_discount", {1}),
+          TensorSpec::createSpec<int32_t>("action_step_type", {1}),
+          TensorSpec::createSpec<float>("action_reward", {1})};
+    } else {
       InputFeatures = {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)};
       TrainingInputFeatures = {
           RA_EVICT_FEATURES_LIST(_DECL_TRAIN_FEATURES)
               TensorSpec::createSpec<float>("action_discount", {1}),
           TensorSpec::createSpec<int32_t>("action_step_type", {1}),
           TensorSpec::createSpec<float>("action_reward", {1})};
     }
+  }
   // support for isa<> and dyn_cast.
   static bool classof(const RegAllocEvictionAdvisorAnalysis *R) {
     return R->getAdvisorMode() == AdvisorMode::Development;
   }
 
   /// get the logger for the given function, or nullptr if we didn't collect
   /// one. This is used to inject the score by the RegAllocScoring pass.
   Logger *getLogger(const MachineFunction &MF) const {
int64_t MLEvictAdvisor::tryFindEvictionCandidatePosition(
   assert(Ret >= 0);
   assert(Ret <= CandidateVirtRegPos);
   return Ret;
 }
 
 bool MLEvictAdvisor::loadInterferenceFeatures(
     const LiveInterval &VirtReg, MCRegister PhysReg, bool IsHint,
     const SmallVirtRegSet &FixedRegisters,
-    llvm::SmallVectorImpl<float> &Largest, size_t Pos) const {
+    llvm::SmallVectorImpl<float> &Largest, size_t Pos,
+    llvm::SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>> &LRPosInfo)
+    const {
   // It is only possible to evict virtual register interference.
   if (Matrix->checkInterference(VirtReg, PhysReg) > LiveRegMatrix::IK_VirtReg) {
     // leave unavailable
     return false;
   }
 
   const bool IsLocal = LIS->intervalIsInOneMBB(VirtReg);
   int64_t LocalIntfs = 0;
for (const LiveInterval *Intf : reverse(IFIntervals)) {
 
       LocalIntfs += (IsLocal && LIS->intervalIsInOneMBB(*Intf) &&
                      (!EnableLocalReassign || !canReassign(*Intf, PhysReg)));
     }
   }
   // OK, so if we made it this far, this LR is an eviction candidate, load its
   // features.
   extractFeatures(InterferingIntervals, Largest, Pos, IsHint, LocalIntfs,
-                  NrUrgent);
+                  NrUrgent, LRPosInfo);
   return true;
 }
 
 MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
     const LiveInterval &VirtReg, const AllocationOrder &Order,
     uint8_t CostPerUseLimit, const SmallVirtRegSet &FixedRegisters) const {
   auto MaybeOrderLimit = getOrderLimit(VirtReg, Order, CostPerUseLimit);
   if (!MaybeOrderLimit)
MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   FeaturesListNormalizer Largest(FeatureIDs::FeatureCount, 0.0);
 
   // Same overal idea as in the default eviction policy - we visit the values
   // of AllocationOrder one at a time. If it's not legally available, we mask
   // off the corresponding feature column (==do nothing because we already
   // reset all the features to 0) Use Pos to capture the column we load
   // features at - in AllocationOrder order.
   size_t Pos = 0;
+  SmallVector<std::tuple<SlotIndex, SlotIndex, size_t>, NumberOfInterferences>
+      LRPosInfo;
   for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit); I != E;
        ++I, ++Pos) {
     MCRegister PhysReg = *I;
     assert(!Regs[Pos].second);
     assert(PhysReg);
     if (!canAllocatePhysReg(CostPerUseLimit, PhysReg)) {
       continue;
     }
     if (loadInterferenceFeatures(VirtReg, PhysReg, I.isHint(), FixedRegisters,
-                                 Largest, Pos)) {
+                                 Largest, Pos, LRPosInfo)) {
       ++Available;
       Regs[Pos] = std::make_pair(PhysReg, true);
     }
   }
   if (Available == 0) {
     // Nothing to decide, nothing to learn.
     assert(!MustFindEviction);
     return MCRegister::NoRegister;
   }
   const size_t ValidPosLimit = Pos;
   // If we must find eviction, the candidate should be masked out of the
   // decision making process.
   Regs[CandidateVirtRegPos].second = !MustFindEviction;
   if (!MustFindEviction)
     extractFeatures(SmallVector<const LiveInterval *, 1>(1, &VirtReg), Largest,
                     CandidateVirtRegPos, /*IsHint*/ 0,
                     /*LocalIntfsCount*/ 0,
-                    /*NrUrgent*/ 0.0);
+                    /*NrUrgent*/ 0.0, LRPosInfo);
   assert(InitialQSize > 0.0 && "We couldn't have gotten here if we had "
                                "nothing to allocate initially.");
+  if (EnableDevelopmentFeatures) {
+    extractInstructionFeatures(LRPosInfo);
+  }
   // Normalize the features.
   for (auto &V : Largest)
     V = V ? V : 1.0;
   for (size_t FeatureIndex = 0; FeatureIndex < FeatureIDs::FeatureCount;
        ++FeatureIndex) {
     if (DoNotNormalize.test(FeatureIndex))
       continue;
     for (size_t Pos = 0; Pos < NumberOfInterferences; ++Pos) {
MLEvictAdvisor::getLIFeatureComponents(const LiveInterval &LI) const {
   return Ret;
 }
 
 // Overall, this currently mimics what we do for weight calculation, but instead
 // of accummulating the various features, we keep them separate.
 void MLEvictAdvisor::extractFeatures(
     const SmallVectorImpl<const LiveInterval *> &Intervals,
     llvm::SmallVectorImpl<float> &Largest, size_t Pos, int64_t IsHint,
-    int64_t LocalIntfsCount, float NrUrgent) const {
+    int64_t LocalIntfsCount, float NrUrgent,
+    SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>> &LRPosInfo)
+    const {
   int64_t NrDefsAndUses = 0;
   int64_t NrBrokenHints = 0;
   double R = 0.0;
   double W = 0.0;
   double RW = 0.0;
   double IndVarUpdates = 0.0;
   double HintWeights = 0.0;
   float StartBBFreq = 0.0;
for (const auto *L : Intervals) {
     R += LIFC.R;
     W += LIFC.W;
     RW += LIFC.RW;
 
     IndVarUpdates += LIFC.IndVarUpdates;
 
     HintWeights += LIFC.HintWeights;
     NrRematerializable += LIFC.IsRemat;
+
+    for (auto CurrentSegment : LI) {
+      LRPosInfo.push_back(
+          std::make_tuple(CurrentSegment.start, CurrentSegment.end, Pos));
+    }
   }
   size_t Size = 0;
   if (!Intervals.empty()) {
     StartBBFreq =
         MBFI.getBlockFreqRelativeToEntryBlock(LIS->getMBBFromIndex(StartSI));
     if (EndSI >= LIS->getSlotIndexes()->getLastIndex())
       EndSI = LIS->getSlotIndexes()->getLastIndex().getPrevIndex();
     EndBBFreq =
#define SET(ID, TYPE, VAL)                                                     \
   SET(hottest_bb_freq_by_max, float, HottestBlockFreq);
   SET(liverange_size, float, Size);
   SET(use_def_density, float, TotalWeight);
   SET(max_stage, int64_t, MaxStage);
   SET(min_stage, int64_t, MinStage);
 #undef SET
 }
 
+void MLEvictAdvisor::extractInstructionFeatures(
+    SmallVectorImpl<std::tuple<SlotIndex, SlotIndex, size_t>> &LRPosInfo)
+    const {
+  std::sort(LRPosInfo.begin(), LRPosInfo.end(),
+            [](std::tuple<SlotIndex, SlotIndex, size_t> A,
+               std::tuple<SlotIndex, SlotIndex, size_t> B) {
+              return std::get<0>(A) < std::get<0>(B);
+            });
+  size_t InstructionCount = 0;
+  size_t CurrentSegment = 0;
+  SlotIndex CurrentIndex = std::get<0>(LRPosInfo[0]);
+  while (true) {
+    while (CurrentIndex <= std::get<1>(LRPosInfo[CurrentSegment]) &&
+           InstructionCount < ModelMaxSupportedInstructionCount) {
+      // set instruction
+      auto *CurrentMachineInstruction =
+          LIS->getInstructionFromIndex(CurrentIndex);
+      if (CurrentMachineInstruction == nullptr) {
+        CurrentIndex = CurrentIndex.getNextIndex();
+        continue;
+      }
+      auto CurrentOpcode = CurrentMachineInstruction->getOpcode();
+      Runner->getTensor<int64_t>(FeatureIDs::instructions -
+                                 1)[InstructionCount] =
+          CurrentOpcode < OpcodeCountCutoff ? CurrentOpcode : 0;
+      // set mask for instruction
+      // add 1 to the resulting position as all of the segment indices are
+      // offset 1 as the first row is instruction opcodes
+      auto CurrentSegmentPosition = std::get<2>(LRPosInfo[CurrentSegment]);
+      Runner->getTensor<int64_t>(
+          FeatureIDs::instructions_mapping -
+          1)[CurrentSegmentPosition * ModelMaxSupportedInstructionCount +
+             InstructionCount] = 1;
+      // handle the overlapping LR case
+      size_t OverlapCheckCurrentSegment = CurrentSegment + 1;
+      while (OverlapCheckCurrentSegment < LRPosInfo.size()) {
+        if (std::get<0>(LRPosInfo[OverlapCheckCurrentSegment]) <=
+            CurrentIndex) {
+          auto OverlapCurrentSegmentPosition =
+              std::get<2>(LRPosInfo[OverlapCheckCurrentSegment]);
+          Runner->getTensor<int64_t>(FeatureIDs::instructions_mapping -
+                                     1)[OverlapCurrentSegmentPosition *
+                                            ModelMaxSupportedInstructionCount +
+                                        InstructionCount] = 1;
+        } else {
+          break;
+        }
+        ++OverlapCheckCurrentSegment;
+      }
+      ++InstructionCount;
+      CurrentIndex = CurrentIndex.getNextIndex();
+    }
+    // if we've just finished processing through the last segment or if we've
+    // hit the maximum number of instructions, break out of the loop.
+    if (CurrentSegment == LRPosInfo.size() - 1 ||
+        InstructionCount >= ModelMaxSupportedInstructionCount) {
+      break;
+    }
+    // just finished processing the previous segment, transition to the next one
+    if (std::get<0>(LRPosInfo[CurrentSegment + 1]) <=
+        std::get<1>(LRPosInfo[CurrentSegment])) {
+      // segments are overlapping.
+      ++CurrentSegment;
+    } else {
+      // segments are not overlapping.
+      CurrentIndex = std::get<0>(LRPosInfo[CurrentSegment + 1]);
+      ++CurrentSegment;
+    }
+  }
+}
+
 // Development mode-specific implementations
 #ifdef LLVM_HAVE_TF_API
 RegAllocEvictionAdvisorAnalysis *llvm::createDevelopmentModeAdvisor() {
   return new DevelopmentModeEvictionAdvisorAnalysis();
 }
 
 int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
     const LiveInterval &VirtReg, const AllocationOrder &Order,
     unsigned OrderLimit, uint8_t CostPerUseLimit,
     const SmallVirtRegSet &FixedRegisters) const {
   int64_t Ret = 0;
   if (isa<ModelUnderTrainingRunner>(getRunner())) {
     Ret = MLEvictAdvisor::tryFindEvictionCandidatePosition(
         VirtReg, Order, OrderLimit, CostPerUseLimit, FixedRegisters);
   } else {
     MCRegister PhysReg = getDefaultAdvisor().tryFindEvictionCandidate(
         VirtReg, Order, CostPerUseLimit, FixedRegisters);
     // Find the index of the selected PhysReg. We need it for logging,
     // otherwise this is wasted cycles (but so would starting development mode
     // without a model nor logging)
     if (!PhysReg)
       Ret = CandidateVirtRegPos;
     else
       for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit);
            I != E; ++I, ++Ret)
         if (*I == PhysReg)
           break;
   }
   if (TrainingLog.empty())
     return Ret;
   size_t CurrentFeature = 0;
   for (; CurrentFeature < FeatureIDs::FeatureCount; ++CurrentFeature) {
     Log->logSpecifiedTensorValue(
         CurrentFeature, reinterpret_cast<const char *>(
                             getRunner().getTensorUntyped(CurrentFeature)));
   }
+  if (EnableDevelopmentFeatures) {
+    // Skip FeatureIDs::FeatureCount
+    for (; CurrentFeature < FeatureIDs::FeaturesWithDevelopmentCount - 1;
+         ++CurrentFeature) {
+      Log->logSpecifiedTensorValue(
+          CurrentFeature, reinterpret_cast<const char *>(
+                              getRunner().getTensorUntyped(CurrentFeature)));
+    }
+  }
   if (auto *MUTR = dyn_cast<ModelUnderTrainingRunner>(&getRunner()))
     for (size_t I = 1; I < MUTR->outputLoggedFeatureSpecs().size();
          ++I, ++CurrentFeature)
       Log->logSpecifiedTensorValue(
           CurrentFeature,
           reinterpret_cast<const char *>(
               MUTR->lastEvaluationResult()->getUntypedTensorValue(I)));
   // The output is right after the features and the extra outputs
   Log->logInt64Value(CurrentFeature, &Ret);
   return Ret;
 }
 
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &MF) {
   if (auto *DevModeAnalysis = dyn_cast<DevelopmentModeEvictionAdvisorAnalysis>(
           &getAnalysis<RegAllocEvictionAdvisorAnalysis>()))
     if (auto *Log = DevModeAnalysis->getLogger(MF))
       Log->logFloatFinalReward(static_cast<float>(
           calculateRegAllocScore(MF, getAnalysis<MachineBlockFrequencyInfo>())
               .getScore()));
 
   return false;
 }
 #endif // #ifdef LLVM_HAVE_TF_API
 
 RegAllocEvictionAdvisorAnalysis *llvm::createReleaseModeAdvisor() {
   return new ReleaseModeEvictionAdvisorAnalysis();
 }
 
 // In all cases except development mode, we don't need scoring.
 #if !defined(LLVM_HAVE_TF_API)
 bool RegAllocScoring::runOnMachineFunction(MachineFunction &) { return false; }
 #endif

llvm/test/CodeGen/MLRegalloc/dev-mode-extra-features-logging.ll

+; REQUIRES: have_tf_api
+; REQUIRES: x86_64-linux
+;
+; Check that we log the currently in development features correctly with both the default
+; case and with a learned policy.
+;
+; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
+; RUN:   -regalloc-training-log=%t1 -tfutils-text-log \
+; RUN:   -regalloc-enable-development-features < %S/Inputs/input.ll
+; RUN: sed -i 's/ \+/ /g' %t1
+; RUN: sed -i 's/\\n key:/\n key:/g' %t1
+; RUN: sed -i 's/\\n feature/\n feature/g' %t1
+; RUN: sed -i 's/\\n/ /g' %t1
+; RUN: FileCheck --input-file %t1 %s
+
+; RUN: rm -rf %t && mkdir %t
+; RUN: %python %S/../../../lib/Analysis/models/gen-regalloc-eviction-test-model.py %t
+; RUN: llc -mtriple=x86_64-linux-unknown -regalloc=greedy -regalloc-enable-advisor=development \
+; RUN:   -regalloc-training-log=%t2 -tfutils-text-log -regalloc-model=%t \
+; RUN:   -regalloc-enable-development-features < %S/Inputs/input.ll
+; RUN: sed -i 's/ \+/ /g' %t2
+; RUN: sed -i 's/\\n key:/\n key:/g' %t2
+; RUN: sed -i 's/\\n feature/\n feature/g' %t2
+; RUN: sed -i 's/\\n/ /g' %t2
+; RUN: FileCheck --input-file %t2 %s
+
+; CHECK-NOT: nan
+; CHECK-LABEL: key: \"instructions\"
+; Check the first five opcodes in the first eviction problem
+; CHECK-NEXT: value: 19
+; CHECK-SAME: value: 19
+; CHECK-SAME: value: 3030
+; CHECK-SAME: value: 1245
+; CHECK-SAME: value: 1264
+; The first eviction problem is significantly less than 300 instructions. Check
+; that there is a zero value
+; CHECK-SAME: value: 0
+; Only the candidate virtreg and the 10th LR are included in this problem. Make
+; sure the other LRs have values of zero.
+; CHECK-LABEL: key: \"instructions_mapping\"
+; CHECK-COUNT-2700: value: 0
+; CHECK-SAME: value: 1
+; Indexing 300 back from where the candidate vr actual resides due to the fact
+; that not all the values between the 10th LR and the candidate are zero.
+; CHECK-COUNT-6600: value: 0
+; CHECK-SAME: value: 1
+; Ensure that we can still go through the mapping matrices for the rest of the
+; eviction problems to make sure we haven't hit the end of the matrix above.
+; There are a total of 23 eviction problems with this test.
+; CHECK-COUNT-22: int64_list
+; CHECK: key: \"is_free\"
 No newline at end of file