[SR-7669] Use SILPrinter callback mechanism to annotate the printed SIL code

[rxwei]

Previous ID	SR-7669
Radar	None
Original Reporter	@rxwei
Type	Bug

Additional Detail from JIRA

Votes	0
Component/s	Swift for TensorFlow
Labels	Bug
Assignee	bgogul (JIRA)
Priority	Medium

md5: 950bb1ab834f9c12bb444d7dc31543f1

Issue Description:

[google@4e27ff2|4e27ff2] added code for annotating SIL code with TFPartition markings.

One follow-up suggestion from clattner is to use sil printer callback before printing each instruction (SILPrinter.cpp:634):

    for (const SILInstruction &I : *BB) {
      Ctx.printInstructionCallBack(&I);
      if (SILPrintGenericSpecializationInfo) {
        if (auto AI = ApplySite::isa(const_cast<SILInstruction *>(&I)))
          if (AI.getSpecializationInfo() && AI.getCalleeFunction())
            printGenericSpecializationInfo(
                PrintState.OS, "call-site", AI.getCalleeFunction()->getName(),
                AI.getSpecializationInfo(), AI.getSubstitutions());
      }
      print(&I);
    }

So we can derive a struct from SILPrintContext, and implement [move] kind of annotations by overriding Ctx.printInstructionCallBack. This way we hopefully just need to call “print” on the function with the right printer context defined.

Example code:

/// A print context which records the line numbers where instructions are
/// printed.
struct AnalysisPrintContext : public SILPrintContext {
  AnalysisPrintContext(llvm::raw_ostream &OS) : SILPrintContext(OS) {}

  ~AnalysisPrintContext() override {}

  void printInstructionCallBack(const SILInstruction *I) override {
    OutStream << "XXXX ";
    I->print(OutStream);
  }
};

Call site:

  if (auto *outs = getTFDumpIntermediateStream()) {
    AnalysisPrintContext ctx(*outs);
    fn.print(ctx);
  }

One challenge is that in addition to annotating instructions, we also want to annotate BB args ([https://github.com/google/swift/blob/4e27ff202817914fe8b776788c8245abdccfae8b/lib/SILOptimizer/Mandatory/TFPartition.cpp#L1877|code here]). So the current callback mechanism on inst level is not sufficient. One option is to extend the SILPrinter callback infra to support annotating BB args.