[ThinLTO] Fix memory corruption in ThinLTOCodeGenerator when CodeGenOnly was specified

[alexbdv]

Summary

Issue occurs when doing ThinLTO with CodeGenOnly flag.
TMBuilder.TheTriple is assigned to by multiple threads in an unsafe way resulting in double-free of std::string memory.

Pseudocode:
if (CodeGenOnly) {

// Perform only parallel codegen and return.
ThreadPool Pool;
int count = 0;
for (auto &ModuleBuffer : Modules) {
  Pool.async([&](int count) { 
  ...
    /// Now call OutputBuffer = codegen(*TheModule);
    /// Which turns into initTMBuilder(moduleTMBuilder, Triple(TheModule.getTargetTriple()));
    /// Which turns into

    TMBuilder.TheTriple = std::move(TheTriple);   // std::string = "....."        
    /// So, basically std::string assignment to same string on multiple threads = memory corruption

}

return;

}

[steven_wu]

I don't think we need another instance of TMBuilder. Can we just call codegenModule directly from run()? TMBuilder should be initialized already from addModule.

diff --git a/lib/LTO/ThinLTOCodeGenerator.cpp b/lib/LTO/ThinLTOCodeGenerator.cpp
index 642e538ecf9..90d4422c796 100644
--- a/lib/LTO/ThinLTOCodeGenerator.cpp
+++ b/lib/LTO/ThinLTOCodeGenerator.cpp
@@ -893,7 +893,7 @@ void ThinLTOCodeGenerator::run() {
                                  /*IsImporting*/ false);

         // CodeGen
-        auto OutputBuffer = codegen(*TheModule);
+        auto OutputBuffer = codegenModule(*TheModule, *TMBuilder.create());
         if (SavedObjectsDirectoryPath.empty())
           ProducedBinaries[count] = std::move(OutputBuffer);
         else

[alexbdv]

Indeed it can, that is another fix. I was just doing the "obvious" fix since I assume there is a reason why the ::codegen(Module &TheModule) function exists.
Should we just remove the function altogether ? As-is it is not thread-safe = results in crashes.
Also, a clone of TMBuilder is not expensive considering the context.

[steven_wu]

I dont mind removing codege(). This interface exists since the first prototype and no longer exposed through C API. It seems to have the same model of optimize(), which is also not thread safe. optimize() is still used by llvm-lto but not in thread context.

[steven_wu]

Remove the declaration in ThinLTOCodeGenerator.h?

[steven_wu]

LGTM. Thanks!

[alexbdv]

I don't have commit access. steven_wu can you please commit this for me ?

llvm/trunk/include/llvm/LTO/legacy/ThinLTOCodeGenerator.h

public:
    */
   void internalize(Module &Module, ModuleSummaryIndex &Index);
 
   /**
    * Perform post-importing ThinLTO optimizations.
    */
   void optimize(Module &Module);
 
-  /**
-   * Perform ThinLTO CodeGen.
-   */
-  std::unique_ptr<MemoryBuffer> codegen(Module &Module);
-
   /**@}*/
 
 private:
   /// Helper factory to build a TargetMachine
   TargetMachineBuilder TMBuilder;
 
   /// Vector holding the in-memory buffer containing the produced binaries, when
   /// SavedObjectsDirectoryPath isn't set.

llvm/trunk/lib/LTO/ThinLTOCodeGenerator.cpp

  */
 void ThinLTOCodeGenerator::optimize(Module &TheModule) {
   initTMBuilder(TMBuilder, Triple(TheModule.getTargetTriple()));
 
   // Optimize now
   optimizeModule(TheModule, *TMBuilder.create(), OptLevel, Freestanding);
 }
 
-/**
- * Perform ThinLTO CodeGen.
- */
-std::unique_ptr<MemoryBuffer> ThinLTOCodeGenerator::codegen(Module &TheModule) {
-  initTMBuilder(TMBuilder, Triple(TheModule.getTargetTriple()));
-  return codegenModule(TheModule, *TMBuilder.create());
-}
-
 /// Write out the generated object file, either from CacheEntryPath or from
 /// OutputBuffer, preferring hard-link when possible.
 /// Returns the path to the generated file in SavedObjectsDirectoryPath.
 static std::string writeGeneratedObject(int count, StringRef CacheEntryPath,
                                         StringRef SavedObjectsDirectoryPath,
                                         const MemoryBuffer &OutputBuffer) {
   SmallString<128> OutputPath(SavedObjectsDirectoryPath);
   llvm::sys::path::append(OutputPath, Twine(count) + ".thinlto.o");
for (auto &ModuleBuffer : Modules) {
         Context.setDiscardValueNames(LTODiscardValueNames);
 
         // Parse module now
         auto TheModule =
             loadModuleFromBuffer(ModuleBuffer.getMemBuffer(), Context, false,
                                  /*IsImporting*/ false);
 
         // CodeGen
-        auto OutputBuffer = codegen(*TheModule);
+        auto OutputBuffer = codegenModule(*TheModule, *TMBuilder.create());
         if (SavedObjectsDirectoryPath.empty())
           ProducedBinaries[count] = std::move(OutputBuffer);
         else
           ProducedBinaryFiles[count] = writeGeneratedObject(
               count, "", SavedObjectsDirectoryPath, *OutputBuffer);
       }, count++);
     }