[BOLT] Attach ORC info to instructions in CFG

[maksfb]

Summary

Propagate Linux Kernel ORC information read from the file to the whole
function CFG once the graph has been built. We have a choice to either
attach ORC state annotation to every instruction, or to the first
instruction in the basic block to conserve processing memory. I chose to
attach to every instruction under --print-orc option which is currently
on by default.

Depends on D155153

[Amir]

bolt/lib/Rewrite/LinuxKernelRewriter.cpp

533

[maksfb]

Fix comment.

bolt/lib/Rewrite/LinuxKernelRewriter.cpp

class LinuxKernelRewriter final : public MetadataRewriter {
   void processLKSMPLocks();
 
   /// Update LKMarkers' locations for the output binary.
   void updateLKMarkers();
 
   /// Read ORC unwind information and annotate instructions.
   Error readORCTables();
 
+  /// Update ORC for functions once CFG is constructed.
+  Error processORCPostCFG();
+
   /// Update ORC data in the binary.
   Error rewriteORCTables();
 
   /// Mark instructions referenced by kernel metadata.
   Error markInstructions();
 
 public:
   LinuxKernelRewriter(BinaryContext &BC)
       : MetadataRewriter("linux-kernel-rewriter", BC) {}
 
   Error preCFGInitializer() override {
     processLKSections();
     if (Error E = markInstructions())
       return E;
 
     if (Error E = readORCTables())
       return E;
 
     return Error::success();
   }
 
+  Error postCFGInitializer() override {
+    if (Error E = processORCPostCFG())
+      return E;
+
+    return Error::success();
+  }
+
   Error postEmitFinalizer() override {
     updateLKMarkers();
 
     if (Error E = rewriteORCTables())
       return E;
 
     return Error::success();
   }
for (uint32_t Index = 0; Index < NumEntries; ++Index) {
     BC.MIB->addAnnotation(*Inst, "ORC", Entry.ORC);
 
     BF->setHasORC(true);
   }
 
   return Error::success();
 }
 
+Error LinuxKernelRewriter::processORCPostCFG() {
+  // Propagate ORC to the rest of the function. We can annotate every
+  // instruction in every function, but to minimize the overhead, we annotate
+  // the first instruction in every basic block to reflect the state at the
+  // entry. This way, the ORC state can be calculated based on annotations
+  // regardless of the basic block layout. Note that if we insert/delete
+  // instructions, we must take care to attach ORC info to the new/deleted ones.
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions())) {
+    if (!BF.hasORC())
+      continue;
+
+    std::optional<ORCState> CurrentState;
+    for (BinaryBasicBlock &BB : BF) {
+      for (MCInst &Inst : BB) {
+        ErrorOr<ORCState> State =
+            BC.MIB->tryGetAnnotationAs<ORCState>(Inst, "ORC");
+
+        if (State) {
+          CurrentState = *State;
+          continue;
+        }
+
+        if (!CurrentState)
+          continue;
+
+        // While printing ORC, attach info to every instruction for convenience.
+        if (opts::PrintORC || &Inst == &BB.front())
+          BC.MIB->addAnnotation(Inst, "ORC", *CurrentState);
+      }
+    }
+  }
+
+  return Error::success();
+}
+
 Error LinuxKernelRewriter::rewriteORCTables() {
   // TODO:
   return Error::success();
 }
 } // namespace
 
 std::unique_ptr<MetadataRewriter>
 llvm::bolt::createLinuxKernelRewriter(BinaryContext &BC) {
   return std::make_unique<LinuxKernelRewriter>(BC);
 }