[mlgo] Add ability to create feature-gated development features in regalloc advisor

[aidengrossman]

Summary

Currently there is no way to add in development features to the ML
regalloc evict advisor which is useful to have when working on feature
engineering/improving the current model. This patch adds in the ability
to add in development features to the ML regalloc evict advisor which
are gated by a runtime flag and not added in at all if not compiled in
LLVM development mode. This sets the stage for future work where we are
planning on upstreaming some of the newer features that we are currently
experimenting with.

[aidengrossman]

Removed debugging artifacts in TFUtil.cpp

[aidengrossman]

This patch adds in a DEVELOPMENT_RA_EVICT_FEATURES_LIST macro that contains development features and modifies/refactors the reset of the file in order to be able run these features when the command line option enable-features-under-development is specified. I tested this in release mode and development mode when compiled in MLGO development mode (setting TENSORFLOW_C_LIB_PATH) and I did another build in release mode (LLVM_HAVE_TF_API is false, TENSORFLOW_C_LIB_PATH isn't set) and everything seemed to work correctly, but it definitely wouldn't hurt to have more people testing this change.

[mtrofin]

Is the main goal of the patch enabling a runtime-determined set of features, to avoid the overhead of allocating buffers for unused features? If yes, I'd just scope it at that change: making the feature storage a vector, and simplifying TrainingInputFeatures as such.

Then I'd just add the change you have in mind, i.e. skip having a 'dummy feature'.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

201

I don't understand how this is used: if the issue is figuring out if we are doing release or development, we can just check the type of the advisor or of the evaluator?

[aidengrossman]

Reduced the scope of the change to just refactoring parts of the
MLRegallocEvictAdvisor to set the stage for future work involving
feature lists that vary at runtime.

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

335

don't init the static in the ctor. Init it outside (i.e. it can be init-ed where it was before)

also, why make it a member? that seems unrelated change

343

no need to make it a member / or separate change.

402

same as above: don't init statics in the ctor, etc.

[aidengrossman]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

335

I've moved the initialization of InputFeatures and TrainingInputFeatures to the constructors of the Analyses that use them because adding extra features at runtime requires the modification of these vectors at runtime. I've kept them as static to share a single copy per analysis across the entire invocation of LLVM (although I don't believe it should matter as if I'm understanding everything correctly, there's only one analysis created per invocation). I've moved them to being class members as I believe it makes reading the code more clear (versus sharing commonalities between the advisors and then extending the vectors in the constructors), but I can change it around if desired. This is necessary for the goal of the patch (refactoring to allow for runtime enabled features) if I'm not mistaken and there isn't some approach I'm unaware of, but I can rescope if desired.

[mtrofin]

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

335

Ack on the initialization, but let's not make them static then. The analysis is shared, indeed.

[aidengrossman]

Changed InputFeatures and TrainingInputFeatures from being static members to
normal class members.

llvm/lib/CodeGen/MLRegallocEvictAdvisor.cpp

#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")
 
-// 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)
+// 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)
 #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.
+// 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
 }
 
-// Per-live interval components that get aggregated into the feature values that
-// will be passed to the evaluator.
+// 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 = std::array<float, FeatureIDs::FeatureCount>;
+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:
   const RegAllocEvictionAdvisor &getDefaultAdvisor() const {
     return static_cast<const RegAllocEvictionAdvisor &>(DefaultAdvisor);
   }
 
   // The assumption is that if the Runner could not be constructed, we emit-ed
   // error, and we shouldn't be asking for it here.
   const MLModelRunner &getRunner() const { return *Runner; }
 
-  /// This just calls Evaluate on the Runner, but in the development mode case,
-  /// if we're just capturing the log of the default advisor, it needs to call
-  /// the latter instead, so we need to pass all the necessary parameters for
-  /// it. In the development case, it will also log.
+  /// This just calls Evaluate on the Runner, but in the development mode
+  /// case, if we're just capturing the log of the default advisor, it needs
+  /// to call the latter instead, so we need to pass all the necessary
+  /// parameters for it. In the development case, it will also log.
   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,
-                           std::array<float, FeatureIDs::FeatureCount> &Largest,
+  bool loadInterferenceFeatures(const LiveInterval &VirtReg, MCRegister PhysReg,
+                                bool IsHint,
+                                const SmallVirtRegSet &FixedRegisters,
+                                llvm::SmallVectorImpl<float> &Largest,
                                 size_t Pos) 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,
-                       std::array<float, FeatureIDs::FeatureCount> &Largest,
-                       size_t Pos, int64_t IsHint, int64_t LocalIntfsCount,
+                       llvm::SmallVectorImpl<float> &Largest, size_t Pos,
+                       int64_t IsHint, int64_t LocalIntfsCount,
                        float NrUrgent) const;
 
-  // Point-in-time: we didn't learn this, so we always delegate to the default.
+  // 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);
   }
 
   const LIFeatureComponents &
   getLIFeatureComponents(const LiveInterval &LI) const;
 
   // Hold on to a default advisor for:
-  // 1) the implementation of canEvictHintInterference, because we didn't learn
-  // that nuance yet;
-  // 2) for bootstrapping (logging) in the development mode case.
+  // 1) the implementation of canEvictHintInterference, because we didn't
+  // learn that nuance yet; 2) for bootstrapping (logging) in the development
+  // mode case.
   const DefaultEvictionAdvisor DefaultAdvisor;
   MLModelRunner *const Runner;
   const MachineBlockFrequencyInfo &MBFI;
   const MachineLoopInfo &Loops;
 
   // Indices of those features we don't want to normalize.
   // This could be static and shared, but its initialization is non-trivial.
   std::bitset<FeatureIDs::FeatureCount> DoNotNormalize;
   const float InitialQSize;
 
   using RegID = unsigned;
   mutable DenseMap<RegID, LIFeatureComponents> CachedFeatures;
 };
 
 #define _DECL_FEATURES(type, name, shape, _)                                   \
   TensorSpec::createSpec<type>(#name, shape),
 
-static const std::vector<TensorSpec> InputFeatures{
-    {RA_EVICT_FEATURES_LIST(_DECL_FEATURES)},
-};
-#undef _DECL_FEATURES
 // ===================================
 // Release (AOT) - specifics
 // ===================================
 class ReleaseModeEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {
 public:
   ReleaseModeEvictionAdvisorAnalysis()
-      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {}
+      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Release) {
+    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;
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   std::unique_ptr<RegAllocEvictionAdvisor>
   getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {

 //
 // Features we log
 #ifdef LLVM_HAVE_TF_API
 static const TensorSpec Output =
     TensorSpec::createSpec<int64_t>(DecisionName, {1});
 static const TensorSpec Reward = TensorSpec::createSpec<float>("reward", {1});
 
 // Features we bind on the model. The tensor names have a prefix, and we also
-// need to include some tensors that are expected to be present by the training
-// algo.
+// need to include some tensors that are expected to be present by the
+// training algo.
 // TODO: can we just get rid of these?
 #define _DECL_TRAIN_FEATURES(type, name, shape, _)                             \
   TensorSpec::createSpec<type>(std::string("action_") + #name, shape),
 
-static const std::vector<TensorSpec> 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})}};
-#undef _DECL_TRAIN_FEATURES
-
 class DevelopmentModeEvictAdvisor : public MLEvictAdvisor {
 public:
   DevelopmentModeEvictAdvisor(const MachineFunction &MF, const RAGreedy &RA,
                               MLModelRunner *Runner,
                               const MachineBlockFrequencyInfo &MBFI,
                               const MachineLoopInfo &Loops, Logger *Log)
       : MLEvictAdvisor(MF, RA, Runner, MBFI, Loops), Log(Log) {}
 
 private:
   int64_t tryFindEvictionCandidatePosition(
       const LiveInterval &VirtReg, const AllocationOrder &Order,
       unsigned OrderLimit, uint8_t CostPerUseLimit,
       const SmallVirtRegSet &FixedRegisters) const override;
 
   Logger *const Log;
 };
 
 class DevelopmentModeEvictionAdvisorAnalysis final
     : public RegAllocEvictionAdvisorAnalysis {
 public:
   DevelopmentModeEvictionAdvisorAnalysis()
-      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Development) {}
+      : RegAllocEvictionAdvisorAnalysis(AdvisorMode::Development) {
+    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 {
     auto I = LogMap.find(MF.getName());
     if (I == LogMap.end())
       return nullptr;
     return I->second.get();
   }
 
 private:
+  std::vector<TensorSpec> InputFeatures;
+  std::vector<TensorSpec> TrainingInputFeatures;
+
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<MachineBlockFrequencyInfo>();
     AU.addRequired<MachineLoopInfo>();
     RegAllocEvictionAdvisorAnalysis::getAnalysisUsage(AU);
   }
 
   // Save all the logs (when requested).
   bool doFinalization(Module &M) override {
getAdvisor(const MachineFunction &MF, const RAGreedy &RA) override {
     return std::make_unique<DevelopmentModeEvictAdvisor>(
         MF, RA, Runner.get(), getAnalysis<MachineBlockFrequencyInfo>(),
         getAnalysis<MachineLoopInfo>(), Log);
   }
 
   std::unique_ptr<MLModelRunner> Runner;
   StringMap<std::unique_ptr<Logger>> LogMap;
 };
+
 #endif //#ifdef LLVM_HAVE_TF_API
 } // namespace
 
 float MLEvictAdvisor::getInitialQueueSize(const MachineFunction &MF) {
   auto &MRI = MF.getRegInfo();
   float Ret = 0.0;
   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
     Register Reg = Register::index2VirtReg(I);
int64_t MLEvictAdvisor::tryFindEvictionCandidatePosition(
   int64_t Ret = Runner->evaluate<int64_t>();
   assert(Ret >= 0);
   assert(Ret <= CandidateVirtRegPos);
   return Ret;
 }
 
 bool MLEvictAdvisor::loadInterferenceFeatures(
     const LiveInterval &VirtReg, MCRegister PhysReg, bool IsHint,
-    const SmallVirtRegSet &FixedRegisters, FeaturesListNormalizer &Largest,
-    size_t Pos) const {
+    const SmallVirtRegSet &FixedRegisters,
+    llvm::SmallVectorImpl<float> &Largest, size_t Pos) 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;
   float NrUrgent = 0.0f;
 
   // The cascade tracking is the same as in the default advisor
   unsigned Cascade = RA.getExtraInfo().getCascadeOrCurrentNext(VirtReg.reg());
 
   SmallVector<const LiveInterval *, MaxInterferences> InterferingIntervals;
   for (MCRegUnitIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
     LiveIntervalUnion::Query &Q = Matrix->query(VirtReg, *Units);
-    // Different from the default heuristic, we don't make any assumptions about
-    // what having more than 10 results in the query may mean.
+    // Different from the default heuristic, we don't make any assumptions
+    // about what having more than 10 results in the query may mean.
     const auto &IFIntervals = Q.interferingVRegs(EvictInterferenceCutoff);
     if (IFIntervals.empty() && InterferingIntervals.empty())
       continue;
     if (IFIntervals.size() >= EvictInterferenceCutoff)
       return false;
     InterferingIntervals.append(IFIntervals.begin(), IFIntervals.end());
     for (const LiveInterval *Intf : reverse(IFIntervals)) {
       assert(Register::isVirtualRegister(Intf->reg()) &&
MCRegister MLEvictAdvisor::tryFindEvictionCandidate(
   if (!MaybeOrderLimit)
     return MCRegister::NoRegister;
   unsigned OrderLimit = *MaybeOrderLimit;
 
   // The heuristic sets initial costs such as, if CostPerUseLimit is
   // max<uint8_t>, then any of the costs of the legally-evictable intervals
   // would be lower. When that happens, one of those will be selected.
   // Therefore, we allow the candidate be selected, unless the candidate is
-  // unspillable, in which case it would be incorrect to not find a register for
-  // it.
+  // unspillable, in which case it would be incorrect to not find a register
+  // for it.
   const bool MustFindEviction =
       (!VirtReg.isSpillable() && CostPerUseLimit == static_cast<uint8_t>(~0u));
   // Number of available candidates - if 0, no need to continue.
   size_t Available = 0;
-  // Make sure we don't have leftover partial state from an attempt where we had
-  // no available candidates and bailed out early.
+  // Make sure we don't have leftover partial state from an attempt where we
+  // had no available candidates and bailed out early.
   resetInputs(*Runner);
 
   // Track the index->register mapping because AllocationOrder doesn't do that
   // and we'd have to scan it.
   // Also track their mask, to write asserts/debug.
   CandidateRegList Regs;
   Regs.fill({0, false});
 
   // Track the largest value of features seen during this eviction session. We
   // only normalize (some of) the float features, but it's just simpler to
   // dimension 'Largest' to all the features, especially since we have the
   // 'DoNotNormalize' list.
-  FeaturesListNormalizer Largest;
-  Largest.fill(0.0);
+  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.
+  // 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;
   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 (Available == 0) {
     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,
+                    CandidateVirtRegPos, /*IsHint*/ 0,
+                    /*LocalIntfsCount*/ 0,
                     /*NrUrgent*/ 0.0);
   assert(InitialQSize > 0.0 && "We couldn't have gotten here if we had "
                                "nothing to allocate initially.");
   // Normalize the features.
   for (auto &V : Largest)
     V = V ? V : 1.0;
   for (size_t FeatureIndex = 0; FeatureIndex < FeatureIDs::FeatureCount;
        ++FeatureIndex) {
Ret.IsRemat = VirtRegAuxInfo::isRematerializable(
       LI, *LIS, *VRM, *MF.getSubtarget().getInstrInfo());
   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,
-    std::array<float, FeatureIDs::FeatureCount> &Largest, size_t Pos,
-    int64_t IsHint, int64_t LocalIntfsCount, float NrUrgent) const {
+    llvm::SmallVectorImpl<float> &Largest, size_t Pos, int64_t IsHint,
+    int64_t LocalIntfsCount, float NrUrgent) 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;
int64_t DevelopmentModeEvictAdvisor::tryFindEvictionCandidatePosition(
     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)
+    // 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;
   }