[MLGO][LoopUnroll] Add logging mechanism for loop partial unroll

[eopXD]

Summary

This is the first part of adding MLGO into LoopUnrollPass. We start
by observing the decision to loop partial unroll.

Loop properties are extracted as input features and the partial unroll
count is the decision. An object MLGOAnalysis will be created in the
first time the LoopPass LoopUnrollPass is triggered, logs will
accumulate and all would be flushed at the end of compilation, in the
dtor of LoopUnrollPass.

Depends on D138368, which introduces LoopPropertiesAnalysis.

[eopXD]

Bump CI.

llvm/include/llvm/Transforms/Scalar/LoopUnrollPass.h

public:
   /// This uses the target information (or flags) to control the thresholds for
   /// different unrolling stategies but supports all of them.
   explicit LoopUnrollPass(LoopUnrollOptions UnrollOpts = {})
       : UnrollOpts(UnrollOpts) {}
 
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
   void printPipeline(raw_ostream &OS,
                      function_ref<StringRef(StringRef)> MapClassName2PassName);
+
+  ~LoopUnrollPass();
 };
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_SCALAR_LOOPUNROLLPASS_H

llvm/lib/Transforms/Scalar/LoopUnrollPass.cpp

 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/LoopAnalysisManager.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/LoopPass.h"
+#include "llvm/Analysis/LoopPropertiesAnalysis.h"
 #include "llvm/Analysis/LoopUnrollAnalyzer.h"
+#include "llvm/Analysis/MLModelRunner.h"
+#include "llvm/Analysis/NoInferenceModelRunner.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/TensorSpec.h"
+#include "llvm/Analysis/Utils/TFUtils.h"
+#include "llvm/Analysis/Utils/TrainingLogger.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Scalar.h"
 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Transforms/Utils.h"
 #include "llvm/Transforms/Utils/LoopPeel.h"
 #include "llvm/Transforms/Utils/LoopSimplify.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/SizeOpts.h"
static cl::opt<unsigned> UnrollThresholdAggressive(
     cl::desc("Threshold (max size of unrolled loop) to use in aggressive (O3) "
              "optimizations"));
 static cl::opt<unsigned>
     UnrollThresholdDefault("unroll-threshold-default", cl::init(150),
                            cl::Hidden,
                            cl::desc("Default threshold (max size of unrolled "
                                     "loop), used in all but O3 optimizations"));
 
+static cl::opt<bool>
+    EnableMLGOUnroll("mlgo-unroll", cl::init(false), cl::Hidden,
+                     cl::desc("Allows MLGO to operate on LoopUnroll"));
+
+static cl::opt<std::string>
+    TrainingLog("mlgo-unroll-training-log", cl::Hidden,
+                cl::desc("Training log for loop partial unroll"));
+
+namespace mlgo_loop_unroll {
+
+/// These features are extracted by LoopPropertiesAnalysis
+/// Tuple of (data_type, variable_name, shape, description)
+#define LOOP_UNROLL_FEATURES_LIST(M)                                           \
+  M(int64_t, loop_size, {1}, "size of loop")                                   \
+  M(int64_t, trip_count, {1}, "static trip count of loop")                     \
+  M(int64_t, is_innermost_loop, {1}, "whether the loop is the innermost loop") \
+  M(int64_t, preheader_blocksize, {1}, "preheader blocksize (by instruction)") \
+  M(int64_t, bb_count, {1}, "number of basic blocks (ignoring subloops)")      \
+  M(int64_t, num_of_loop_latch, {1}, "number of loop latches")                 \
+  M(int64_t, load_inst_count, {1}, "load instruction count")                   \
+  M(int64_t, store_inst_count, {1}, "store instruction count")                 \
+  M(int64_t, logical_inst_count, {1}, "logical instruction count")             \
+  M(int64_t, cast_inst_count, {1}, "cast instruction count")
+
+// The model learns to decide whether or not to partial unroll a loop.
+// If unroll_count == 0, a loop is not unrolled, otherwise it is unrolled by the
+// factor of provided decision.
+#define DecisionName "unroll_count"
+
+enum FeatureIDs {
+#define _FEATURE_IDX(_, name, __, ___) name,
+  LOOP_UNROLL_FEATURES_LIST(_FEATURE_IDX)
+#undef _FEATURE_IDX
+      FeatureCount
+};
+
+#define DECL_FEATURES(type, name, shape, _)                                    \
+  TensorSpec::createSpec<type>(#name, shape),
+
+/// Input features for training. Note that in the future we can attempt to have
+/// multiple sets of features for different purpose. For example a set of
+/// input features in release mode.
+static const std::vector<TensorSpec> InputFeatures{
+    LOOP_UNROLL_FEATURES_LIST(DECL_FEATURES)};
+
+/// A single output feature, the decision whether or not to loop partial unroll
+static const TensorSpec OutputFeature =
+    TensorSpec::createSpec<int32_t>("unroll_count", {1});
+
+/// Currently a dummy reward at the moment. Required to be provided for
+/// construction of logger.
+static const TensorSpec RewardFeature =
+    TensorSpec::createSpec<float>("reward", {1});
+
+/// Class for MLGO in loop unroll. Since LoopUnrollPass is a LoopPass and we
+/// will flush the logs after all features all collected, the compiler will
+/// hold a local copy of this class. Logs will be dumped at the end of
+/// compilation, which will be at the dtor of LoopUnrollPass.
+struct MLGOLoopUnrollAnalysis {
+public:
+  /// Ctor for MLGO in loop unroll.
+  MLGOLoopUnrollAnalysis(LLVMContext &Ctx) : Ctx(Ctx) {
+    Runner = std::make_unique<NoInferenceModelRunner>(Ctx, InputFeatures);
+  }
+
+  MLGOLoopUnrollAnalysis() = delete;
+
+  /// Set features for loop
+  void setFeatures(const unsigned LoopSize, const unsigned TripCount,
+                   LoopInfo &LI, ScalarEvolution &SE, Loop &L);
+
+  /// Log features and partial unroll decision for loop
+  void logFeaturesAndDecision(const unsigned PartialUnrollCount, Loop &L);
+
+  /// Save **ALL** the logs
+  bool flush();
+
+  /// Runner for MLModel (training(default) / training(model) / release)
+  std::unique_ptr<MLModelRunner> Runner;
+
+private:
+  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() {
+#define _RESET(type, name, shape, _)                                           \
+  std::memset(Runner->getTensorUntyped(FeatureIDs::name), 0,                   \
+              getTotalSize<type>(shape));
+    LOOP_UNROLL_FEATURES_LIST(_RESET)
+#undef _RESET
+  }
+
+  LLVMContext &Ctx;
+
+  /// A logger for each loop. Key = "$(MODULE)###$(FUNCTION)###$(LOOP)"
+  StringMap<std::unique_ptr<Logger>> LogMap;
+};
+
+void MLGOLoopUnrollAnalysis::setFeatures(const unsigned LoopSize,
+                                         const unsigned TripCount, LoopInfo &LI,
+                                         ScalarEvolution &SE, Loop &L) {
+  resetInputs();
+
+  LoopPropertiesInfo LPI =
+      LoopPropertiesInfo::getLoopPropertiesInfo(&L, &LI, &SE);
+
+#define SET(id, type, val)                                                     \
+  *Runner->getTensor<type>(FeatureIDs::id) = static_cast<type>(val);
+  SET(loop_size, int64_t, LoopSize);
+  SET(trip_count, int64_t, TripCount);
+  SET(is_innermost_loop, int64_t, LPI.IsInnerMostLoop);
+  SET(preheader_blocksize, int64_t, LPI.PreheaderBlocksize);
+  SET(bb_count, int64_t, LPI.BasicBlockCount);
+  SET(num_of_loop_latch, int64_t, LPI.LoopLatchCount);
+  SET(load_inst_count, int64_t, LPI.LoadInstCount);
+  SET(store_inst_count, int64_t, LPI.StoreInstCount);
+  SET(logical_inst_count, int64_t, LPI.LogicalInstCount);
+  SET(cast_inst_count, int64_t, LPI.CastInstCount);
+#undef SET
+}
+
+void MLGOLoopUnrollAnalysis::logFeaturesAndDecision(const unsigned UnrollCount,
+                                                    Loop &L) {
+  // Key = $(MODULE)###$(FUNCTION)###$(LOOP)
+  std::string Key = L.getHeader()->getModule()->getName().str() + "###" +
+                    L.getHeader()->getParent()->getName().str() + "###" +
+                    L.getName().str();
+
+  assert(!LogMap.count(Key) &&
+         "Should only extract feature for every loop once");
+
+  std::vector<LoggedFeatureSpec> LFS;
+  for (const auto &IF : InputFeatures)
+    LFS.push_back({IF, None});
+  LFS.push_back({OutputFeature, None});
+
+  // Create Logger for loop and insert it to LogMap
+  auto I = LogMap.insert((std::make_pair(
+      Key, std::make_unique<Logger>(LFS, RewardFeature,
+                                    /* IncludeReward */ false))));
+  assert(I.second && "Should be unique insertion");
+
+  Logger *Log = I.first->second.get();
+  size_t CurrentFeature = 0;
+  for (; CurrentFeature < FeatureIDs::FeatureCount; ++CurrentFeature)
+    Log->logSpecifiedTensorValue(CurrentFeature,
+                                 reinterpret_cast<const char *>(
+                                     Runner->getTensorUntyped(CurrentFeature)));
+
+  Log->logInt32Value(CurrentFeature,
+                     reinterpret_cast<const int32_t *>(&UnrollCount));
+
+  LLVM_DEBUG(
+      dbgs() << "(MLGO) Logged features and loop partial unroll decision = "
+             << UnrollCount << " for loop '" << Key << "'\n");
+}
+
+bool MLGOLoopUnrollAnalysis::flush() {
+  std::error_code EC;
+  // The append flag is specified here because modules is compiled at the same
+  // time during usual compulations and we want the logs to be in the same file.
+  auto OS = std::make_unique<raw_fd_ostream>(TrainingLog, EC,
+                                             llvm::sys::fs::OF_Append);
+  if (EC) {
+    Ctx.emitError(EC.message() + ":" + TrainingLog);
+    return false;
+  }
+  Logger::flushLogs(*OS, LogMap);
+  return true;
+}
+
+} // namespace mlgo_loop_unroll
+
+static mlgo_loop_unroll::MLGOLoopUnrollAnalysis *MLGOAnalysis = nullptr;
+
 /// A magic value for use with the Threshold parameter to indicate
 /// that the loop unroll should be performed regardless of how much
 /// code expansion would result.
 static const unsigned NoThreshold = std::numeric_limits<unsigned>::max();
 
 /// Gather the various unrolling parameters based on the defaults, compiler
 /// flags, TTI overrides and user specified parameters.
 TargetTransformInfo::UnrollingPreferences llvm::gatherUnrollingPreferences(
if (Cost->UnrolledCost < UP.Threshold * Boost / 100)
       return FullUnrollTripCount;
   }
   return None;
 }
 
 static Optional<unsigned>
 shouldPartialUnroll(const unsigned LoopSize, const unsigned TripCount,
                     const UnrollCostEstimator UCE,
-                    const TargetTransformInfo::UnrollingPreferences &UP) {
+                    const TargetTransformInfo::UnrollingPreferences &UP,
+                    ScalarEvolution &SE, LoopInfo &LI, Loop &L) {
 
   if (!TripCount)
     return None;
 
   if (!UP.Partial) {
     LLVM_DEBUG(dbgs() << "  will not try to unroll partially because "
                << "-unroll-allow-partial not given\n");
     return 0;
   }
+
+  if (EnableMLGOUnroll)
+    MLGOAnalysis->setFeatures(LoopSize, TripCount, LI, SE, L);
+
   unsigned count = UP.Count;
   if (count == 0)
     count = TripCount;
   if (UP.PartialThreshold != NoThreshold) {
     // Reduce unroll count to be modulo of TripCount for partial unrolling.
     if (UCE.getUnrolledLoopSize(UP, count) > UP.PartialThreshold)
       count = (std::max(UP.PartialThreshold, UP.BEInsns + 1) - UP.BEInsns) /
         (LoopSize - UP.BEInsns);
shouldPartialUnroll(const unsigned LoopSize, const unsigned TripCount,
   } else {
     count = TripCount;
   }
   if (count > UP.MaxCount)
     count = UP.MaxCount;
 
   LLVM_DEBUG(dbgs() << "  partially unrolling with count: " << count << "\n");
 
+  if (EnableMLGOUnroll)
+    MLGOAnalysis->logFeaturesAndDecision(count, L);
+
   return count;
 }
 // Returns true if unroll count was set explicitly.
 // Calculates unroll count and writes it to UP.Count.
 // Unless IgnoreUser is true, will also use metadata and command-line options
 // that are specific to to the LoopUnroll pass (which, for instance, are
 // irrelevant for the LoopUnrollAndJam pass).
 // FIXME: This function is used by LoopUnroll and LoopUnrollAndJam, but consumes
bool llvm::computeUnrollCount(
 
   // Before starting partial unrolling, set up.partial to true,
   // if user explicitly asked  for unrolling
   if (TripCount)
     UP.Partial |= ExplicitUnroll;
 
   // 6th priority is partial unrolling.
   // Try partial unroll only when TripCount could be statically calculated.
-  if (auto UnrollFactor = shouldPartialUnroll(LoopSize, TripCount, UCE, UP)) {
+  if (auto UnrollFactor =
+          shouldPartialUnroll(LoopSize, TripCount, UCE, UP, SE, *LI, *L)) {
     UP.Count = *UnrollFactor;
 
     if ((PragmaFullUnroll || PragmaEnableUnroll) && TripCount &&
         UP.Count != TripCount)
       ORE->emit([&]() {
         return OptimizationRemarkMissed(DEBUG_TYPE,
                                         "FullUnrollAsDirectedTooLarge",
                                         L->getStartLoc(), L->getHeader())
if (auto *LAMProxy = AM.getCachedResult<LoopAnalysisManagerFunctionProxy>(F))
     LAM = &LAMProxy->getManager();
 
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
   ProfileSummaryInfo *PSI =
       MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
   auto *BFI = (PSI && PSI->hasProfileSummary()) ?
       &AM.getResult<BlockFrequencyAnalysis>(F) : nullptr;
 
+  // Creating the MLGOAnalysis object at first sight
+  if (EnableMLGOUnroll && !MLGOAnalysis) {
+    LLVM_DEBUG(dbgs() << "(MLGO) Create MLGOAnalysis for loop unroll\n");
+    MLGOAnalysis =
+        new mlgo_loop_unroll::MLGOLoopUnrollAnalysis(SE.getContext());
+  }
+
   bool Changed = false;
 
   // The unroller requires loops to be in simplified form, and also needs LCSSA.
   // Since simplification may add new inner loops, it has to run before the
   // legality and profitability checks. This means running the loop unroller
   // will simplify all loops, regardless of whether anything end up being
   // unrolled.
   for (const auto &L : LI) {
void LoopUnrollPass::printPipeline(
   if (UnrollOpts.AllowProfileBasedPeeling != None)
     OS << (UnrollOpts.AllowProfileBasedPeeling.value() ? "" : "no-")
        << "profile-peeling;";
   if (UnrollOpts.FullUnrollMaxCount != None)
     OS << "full-unroll-max=" << UnrollOpts.FullUnrollMaxCount << ";";
   OS << "O" << UnrollOpts.OptLevel;
   OS << ">";
 }
+
+LoopUnrollPass::~LoopUnrollPass() {
+  // At end of compilation, the LoopUnrollPass dtor is triggered. All logs of
+  // loop features and partial unroll decisions will be collected at this point
+  // and we now dump the logs.
+  if (EnableMLGOUnroll)
+    MLGOAnalysis->flush();
+}