Merge pull request #29326 from eeckstein/irgen-for-merge-functions

IRGen: in multi-threaded compilation, create all specializations of a function in the same LLVM module.

Author: eeckstein

include/swift/SIL/SILFunction.h

@@ -863,6 +863,10 @@ class SILFunction
     validateSubclassScope(getClassSubclassScope(), isThunk(), Info);
     SpecializationInfo = Info;
   }
+  
+  /// If this function is a specialization, return the original function from
+  /// which this function was specialized.
+  const SILFunction *getOriginOfSpecialization() const;
 
   /// Retrieve the generic environment containing the mapping from interface
   /// types to context archetypes for this function. Only present if the

lib/IRGen/GenDecl.cpp

@@ -1209,6 +1209,16 @@ void IRGenerator::addLazyFunction(SILFunction *f) {
   assert(!FinishedEmittingLazyDefinitions);
   LazyFunctionDefinitions.push_back(f);
 
+  if (const SILFunction *orig = f->getOriginOfSpecialization()) {
+    // f is a specialization. Try to emit all specializations of the same
+    // original function into the same IGM. This increases the chances that
+    // specializations are merged by LLVM's function merging.
+    auto iter =
+      IGMForSpecializations.insert(std::make_pair(orig, CurrentIGM)).first;
+    DefaultIGMForFunction.insert(std::make_pair(f, iter->second));
+    return;
+  }
+
   if (auto *dc = f->getDeclContext())
     if (dc->getParentSourceFile())
       return;

lib/IRGen/IRGenModule.h

@@ -204,7 +204,13 @@ class IRGenerator {
   // Stores the IGM from which a function is referenced the first time.
   // It is used if a function has no source-file association.
   llvm::DenseMap<SILFunction *, IRGenModule *> DefaultIGMForFunction;
-  
+
+  // The IGMs where sepecializations of functions are emitted. The key is the
+  // non-specialized function.
+  // Storing all specializations of a function in the same IGM increases the
+  // chances of function merging.
+  llvm::DenseMap<const SILFunction *, IRGenModule *> IGMForSpecializations;
+
   // The IGM of the first source file.
   IRGenModule *PrimaryIGM = nullptr;
 

lib/SIL/SILFunction.cpp

@@ -167,6 +167,17 @@ bool SILFunction::hasForeignBody() const {
   return SILDeclRef::isClangGenerated(getClangNode());
 }
 
+const SILFunction *SILFunction::getOriginOfSpecialization() const {
+  if (!isSpecialization())
+    return nullptr;
+
+  const SILFunction *p = getSpecializationInfo()->getParent();
+  while (p->isSpecialization()) {
+    p = p->getSpecializationInfo()->getParent();
+  }
+  return p;
+}
+
 void SILFunction::numberValues(llvm::DenseMap<const SILNode*, unsigned> &
                                  ValueToNumberMap) const {
   unsigned idx = 0;

test/IRGen/Inputs/multithread_module/main.swift

@@ -23,6 +23,10 @@ public struct MyStruct : MyProto {
 	return x + 3
 }
 
+public func mutateMyStructArray(_ arr: inout [MyStruct], _ x: MyStruct) {
+  arr.append(x)
+}
+
 public var g1 = 234
 
 let i = testit(27)

test/IRGen/multithread_module.swift

@@ -58,8 +58,18 @@ func callproto(_ p: MyProto) {
 	print(p.protofunc())
 }
 
+public func mutateBaseArray(_ arr: inout [Base], _ x: Base) {
+  arr.append(x)
+}
+
+
 // Check the llvm IR files:
 
+// Check if all specializations from stdlib functions are created in the same LLVM module.
+
+// CHECK-MAINLL-DAG: define {{.*}} @"$sSa16_createNewBuffer14bufferIsUnique15minimumCapacity13growForAppendySb_SiSbtF4test8MyStructV_Tg5"
+// CHECK-MAINLL-DAG: define {{.*}} @"$sSa16_createNewBuffer14bufferIsUnique15minimumCapacity13growForAppendySb_SiSbtF4test4BaseC_Tg5"
+
 // Check if the DI filename is correct and not "<unknown>".
 
 // CHECK-MAINLL: [[F:![0-9]+]] = !DIFile(filename: "{{.*}}IRGen/Inputs/multithread_module/main.swift", directory: "{{.*}}")