[JITLink][RISCV] Move relax to PostAllocationPasses

[jobnoorman]

Summary

JITLinkContext is notified (using notifyResolved) of the final
symbol addresses after allocating memory and running the post-allocation
passes. However, linker relaxation, which can cause symbol addresses to
change, was run during the pre-fixup passes. This causes users of
JITLink (e.g., ORC) to pick-up wrong symbol addresses when linker
relaxation was enabled.

This patch fixes this by running relaxation during the post-allocation
passes.

Fixes #63671

@lhames: afaict, llvm-jitlink gathers symbol values during a custom
post-fixup pass [1]. Therefore, it isn't affected by this bug and I'm
not quite sure how to add a test for it. Any suggestions?

[1] https://github.com/llvm/llvm-project/blob/3003da71540bb8483615e18f8ab379da39fa4512/llvm/tools/llvm-jitlink/llvm-jitlink.cpp#L1080-L1082

[Hahnfeld]

Perfect, solves my issue - thanks for the quick fix!

[lhames]

LGTM.

afaict, llvm-jitlink gathers symbol values during a custom
post-fixup pass [1]. Therefore, it isn't affected by this bug and I'm
not quite sure how to add a test for it. Any suggestions?

One approach would be to move the anonymous pass that calls register<Format>GraphInfo in llvm-jitlink.cpp to the post-allocation phase. If I do that I can immediately see the relaxation testcases failing. That doesn't work in general though because later phases (pre- and post-fixup) might introduce new symbols that llvm-jitlink needs to know about -- I see other testcases breaking because of this.

Maybe we could add a post-allocation pass in llvm-jitlink that records the initial addresses, then have register<Format>GraphInfo throw an error if any of the previously resolved symbols has changed address.

It'd be even better if we could catch this in an assert in LinkGraph itself. If we modified the Section class to have a pointer back to the LinkGraph and modified the LinkGraph to track the phase, then we could write checks like:

void Addressable::setAddress(ExecutorAddr Adds) {
  assert(
    (!isDefined() || static_cast<Block&>(*this).getSection().getGraph().getPhase() < PreFixupPass) &&
    "Cannot change defined symbol address after symbol resolution");
}

I think you can land this without waiting for that though. :)

[jobnoorman]

In D154501#4474483, @lhames wrote:

It'd be even better if we could catch this in an assert in LinkGraph itself. If we modified the Section class to have a pointer back to the LinkGraph and modified the LinkGraph to track the phase, then we could write checks like:

void Addressable::setAddress(ExecutorAddr Adds) {
  assert(
    (!isDefined() || static_cast<Block&>(*this).getSection().getGraph().getPhase() < PreFixupPass) &&
    "Cannot change defined symbol address after symbol resolution");
}

This looks like an interesting approach but just a heads-up that if you were to implement this now, it might break BOLT. We currently reassign section addresses within notifyResolved. I know this is not the correct approach and I have a local fix for it. This is part of a bigger patch though and it hasn't been submitted yet.

llvm/include/llvm/ExecutionEngine/JITLink/ELF_riscv.h

 Expected<std::unique_ptr<LinkGraph>>
 createLinkGraphFromELFObject_riscv(MemoryBufferRef ObjectBuffer);
 
 /// jit-link the given object buffer, which must be a ELF riscv object file.
 void link_ELF_riscv(std::unique_ptr<LinkGraph> G,
                     std::unique_ptr<JITLinkContext> Ctx);
 
 /// Returns a pass that performs linker relaxation. Should be added to
-/// PreFixupPasses.
+/// PostAllocationPasses.
 LinkGraphPassFunction createRelaxationPass_ELF_riscv();
 
 } // end namespace jitlink
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_JITLINK_ELF_RISCV_H

llvm/lib/ExecutionEngine/JITLink/ELF_riscv.cpp

void link_ELF_riscv(std::unique_ptr<LinkGraph> G,
   const Triple &TT = G->getTargetTriple();
   if (Ctx->shouldAddDefaultTargetPasses(TT)) {
     if (auto MarkLive = Ctx->getMarkLivePass(TT))
       Config.PrePrunePasses.push_back(std::move(MarkLive));
     else
       Config.PrePrunePasses.push_back(markAllSymbolsLive);
     Config.PostPrunePasses.push_back(
         PerGraphGOTAndPLTStubsBuilder_ELF_riscv::asPass);
-    Config.PreFixupPasses.push_back(relax);
+    Config.PostAllocationPasses.push_back(relax);
   }
   if (auto Err = Ctx->modifyPassConfig(*G, Config))
     return Ctx->notifyFailed(std::move(Err));
 
   ELFJITLinker_riscv::link(std::move(Ctx), std::move(G), std::move(Config));
 }
 
 LinkGraphPassFunction createRelaxationPass_ELF_riscv() { return relax; }
 
 } // namespace jitlink
 } // namespace llvm