Merge pull request #66844 from kubamracek/static-init-structs

Allow using structs with trivial initializers in globals that require static initialization (e.g. @_section attribute)

Author: kubamracek

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/SimplificationPasses.swift

@@ -37,7 +37,7 @@ protocol LateOnoneSimplifyable : Instruction {
 let ononeSimplificationPass = FunctionPass(name: "onone-simplification") {
   (function: Function, context: FunctionPassContext) in
 
-  runSimplification(on: function, context, preserveDebugInfo: true) {
+  _ = runSimplification(on: function, context, preserveDebugInfo: true) {
     if let i = $0 as? OnoneSimplifyable {
       i.simplify($1)
     }
@@ -47,7 +47,7 @@ let ononeSimplificationPass = FunctionPass(name: "onone-simplification") {
 let simplificationPass = FunctionPass(name: "simplification") {
   (function: Function, context: FunctionPassContext) in
 
-  runSimplification(on: function, context, preserveDebugInfo: false) {
+  _ = runSimplification(on: function, context, preserveDebugInfo: false) {
     if let i = $0 as? Simplifyable {
       i.simplify($1)
     }
@@ -57,7 +57,7 @@ let simplificationPass = FunctionPass(name: "simplification") {
 let lateOnoneSimplificationPass = FunctionPass(name: "late-onone-simplification") {
   (function: Function, context: FunctionPassContext) in
 
-  runSimplification(on: function, context, preserveDebugInfo: true) {
+  _ = runSimplification(on: function, context, preserveDebugInfo: true) {
     if let i = $0 as? LateOnoneSimplifyable {
       i.simplifyLate($1)
     } else if let i = $0 as? OnoneSimplifyable {
@@ -73,13 +73,15 @@ let lateOnoneSimplificationPass = FunctionPass(name: "late-onone-simplification"
 
 func runSimplification(on function: Function, _ context: FunctionPassContext,
                        preserveDebugInfo: Bool,
-                       _ simplify: (Instruction, SimplifyContext) -> ()) {
+                       _ simplify: (Instruction, SimplifyContext) -> ()) -> Bool {
   var worklist = InstructionWorklist(context)
   defer { worklist.deinitialize() }
 
+  var changed = false
   let simplifyCtxt = context.createSimplifyContext(preserveDebugInfo: preserveDebugInfo,
                                                    notifyInstructionChanged: {
     worklist.pushIfNotVisited($0)
+    changed = true
   })
 
   // Push in reverse order so that popping from the tail of the worklist visits instruction in forward order again.
@@ -97,7 +99,7 @@ func runSimplification(on function: Function, _ context: FunctionPassContext,
         continue
       }
       if !context.continueWithNextSubpassRun(for: instruction) {
-        return
+        return changed
       }
       simplify(instruction, simplifyCtxt)
     }
@@ -110,6 +112,8 @@ func runSimplification(on function: Function, _ context: FunctionPassContext,
   if context.needFixStackNesting {
     function.fixStackNesting(context)
   }
+  
+  return changed
 }
 
 private func cleanupDeadInstructions(in function: Function,

SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift

@@ -12,7 +12,8 @@
 
 import SIL
 
-/// Performs mandatory optimizations for performance-annotated functions.
+/// Performs mandatory optimizations for performance-annotated functions, and global
+/// variable initializers that are required to be statically initialized.
 ///
 /// Optimizations include:
 /// * de-virtualization
@@ -22,14 +23,15 @@ import SIL
 /// * dead alloc elimination
 /// * instruction simplification
 ///
-/// The pass starts with performance-annotated functions and transitively handles
+/// The pass starts with performance-annotated functions / globals and transitively handles
 /// called functions.
 ///
 let mandatoryPerformanceOptimizations = ModulePass(name: "mandatory-performance-optimizations") {
   (moduleContext: ModulePassContext) in
 
   var worklist = FunctionWorklist()
   worklist.addAllPerformanceAnnotatedFunctions(of: moduleContext)
+  worklist.addAllAnnotatedGlobalInitOnceFunctions(of: moduleContext)
 
   optimizeFunctionsTopDown(using: &worklist, moduleContext)
 }
@@ -47,34 +49,43 @@ private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
   }
 }
 
+fileprivate struct PathFunctionTuple: Hashable {
+  var path: SmallProjectionPath
+  var function: Function
+}
+
 private func optimize(function: Function, _ context: FunctionPassContext) {
-  runSimplification(on: function, context, preserveDebugInfo: true) { instruction, simplifyCtxt in
-    if let i = instruction as? OnoneSimplifyable {
-      i.simplify(simplifyCtxt)
-      if instruction.isDeleted {
-        return
+  var alreadyInlinedFunctions: Set<PathFunctionTuple> = Set()
+  
+  var changed = true
+  while changed {
+    changed = runSimplification(on: function, context, preserveDebugInfo: true) { instruction, simplifyCtxt in
+      if let i = instruction as? OnoneSimplifyable {
+        i.simplify(simplifyCtxt)
+        if instruction.isDeleted {
+          return
+        }
+      }
+      switch instruction {
+      case let apply as FullApplySite:
+        inlineAndDevirtualize(apply: apply, alreadyInlinedFunctions: &alreadyInlinedFunctions, context, simplifyCtxt)
+      default:
+        break
       }
     }
-    switch instruction {
-    case let apply as FullApplySite:
-      inlineAndDevirtualize(apply: apply, context, simplifyCtxt)
-    default:
-      break
-    }
-  }
 
-  _ = context.specializeApplies(in: function, isMandatory: true)
+    _ = context.specializeApplies(in: function, isMandatory: true)
 
-  removeUnusedMetatypeInstructions(in: function, context)
+    removeUnusedMetatypeInstructions(in: function, context)
 
-  // If this is a just specialized function, try to optimize copy_addr, etc.
-  if context.optimizeMemoryAccesses(in: function) {
+    // If this is a just specialized function, try to optimize copy_addr, etc.
+    changed = context.optimizeMemoryAccesses(in: function) || changed
     _ = context.eliminateDeadAllocations(in: function)
   }
 }
 
-private func inlineAndDevirtualize(apply: FullApplySite, _ context: FunctionPassContext, _ simplifyCtxt: SimplifyContext) {
-
+private func inlineAndDevirtualize(apply: FullApplySite, alreadyInlinedFunctions: inout Set<PathFunctionTuple>,
+                                   _ context: FunctionPassContext, _ simplifyCtxt: SimplifyContext) {
   if simplifyCtxt.tryDevirtualize(apply: apply, isMandatory: true) != nil {
     return
   }
@@ -88,7 +99,7 @@ private func inlineAndDevirtualize(apply: FullApplySite, _ context: FunctionPass
     return
   }
 
-  if shouldInline(apply: apply, callee: callee) {
+  if shouldInline(apply: apply, callee: callee, alreadyInlinedFunctions: &alreadyInlinedFunctions) {
     simplifyCtxt.inlineFunction(apply: apply, mandatoryInline: true)
 
     // In OSSA `partial_apply [on_stack]`s are represented as owned values rather than stack locations.
@@ -110,7 +121,7 @@ private func removeUnusedMetatypeInstructions(in function: Function, _ context:
   }
 }
 
-private func shouldInline(apply: FullApplySite, callee: Function) -> Bool {
+private func shouldInline(apply: FullApplySite, callee: Function, alreadyInlinedFunctions: inout Set<PathFunctionTuple>) -> Bool {
   if callee.isTransparent {
     return true
   }
@@ -123,9 +134,103 @@ private func shouldInline(apply: FullApplySite, callee: Function) -> Bool {
     // Force inlining them in global initializers so that it's possible to statically initialize the global.
     return true
   }
+  if apply.parentFunction.isGlobalInitOnceFunction,
+      let global = apply.parentFunction.getInitializedGlobal(),
+      global.mustBeInitializedStatically,
+      let applyInst = apply as? ApplyInst,
+      let projectionPath = applyInst.isStored(to: global),
+      alreadyInlinedFunctions.insert(PathFunctionTuple(path: projectionPath, function: callee)).inserted {
+    return true
+  }
   return false
 }
 
+private extension Value {
+  /// Analyzes the def-use chain of an apply instruction, and looks for a single chain that leads to a store instruction
+  /// that initializes a part of a global variable or the entire variable:
+  ///
+  /// Example:
+  ///   %g = global_addr @global
+  ///   ...
+  ///   %f = function_ref @func
+  ///   %apply = apply %f(...)
+  ///   store %apply to %g   <--- is a store to the global trivially (the apply result is immediately going into a store)
+  ///
+  /// Example:
+  ///   %apply = apply %f(...)
+  ///   %apply2 = apply %f2(%apply)
+  ///   store %apply2 to %g   <--- is a store to the global (the apply result has a single chain into the store)
+  ///
+  /// Example:
+  ///   %a = apply %f(...)
+  ///   %s = struct $MyStruct (%a, %b)
+  ///   store %s to %g   <--- is a partial store to the global (returned SmallProjectionPath is MyStruct.s0)
+  ///
+  /// Example:
+  ///   %a = apply %f(...)
+  ///   %as = struct $AStruct (%other, %a)
+  ///   %bs = struct $BStruct (%as, %bother)
+  ///   store %bs to %g   <--- is a partial store to the global (returned SmallProjectionPath is MyStruct.s0.s1)
+  ///
+  /// Returns nil if we cannot find a singular def-use use chain (e.g. because a value has more than one user)
+  /// leading to a store to the specified global variable.
+  func isStored(to global: GlobalVariable) -> SmallProjectionPath? {
+    var singleUseValue: any Value = self
+    var path = SmallProjectionPath()
+    while true {
+      guard let use = singleUseValue.uses.singleNonDebugUse else {
+        return nil
+      }
+      
+      switch use.instruction {
+      case is StructInst:
+        path = path.push(.structField, index: use.index)
+        break
+      case is TupleInst:
+        path = path.push(.tupleField, index: use.index)
+        break
+      case let ei as EnumInst:
+        path = path.push(.enumCase, index: ei.caseIndex)
+        break
+      case let si as StoreInst:
+        guard let storeDestination = si.destination as? GlobalAddrInst else {
+          return nil
+        }
+
+        guard storeDestination.global == global else {
+          return nil
+        }
+        
+        return path
+      default:
+        return nil
+      }
+
+      guard let nextInstruction = use.instruction as? SingleValueInstruction else {
+        return nil
+      }
+
+      singleUseValue = nextInstruction
+    }
+  }
+}
+
+private extension Function {
+  /// Analyzes the global initializer function and returns global it initializes (from `alloc_global` instruction).
+  func getInitializedGlobal() -> GlobalVariable? {
+    for inst in self.entryBlock.instructions {
+      switch inst {
+      case let agi as AllocGlobalInst:
+        return agi.global
+      default:
+        break
+      }
+    }
+
+    return nil
+  }
+}
+
 fileprivate struct FunctionWorklist {
   private(set) var functions = Array<Function>()
   private var pushedFunctions = Set<Function>()
@@ -146,6 +251,15 @@ fileprivate struct FunctionWorklist {
     }
   }
 
+  mutating func addAllAnnotatedGlobalInitOnceFunctions(of moduleContext: ModulePassContext) {
+    for f in moduleContext.functions where f.isGlobalInitOnceFunction {
+      if let global = f.getInitializedGlobal(),
+         global.mustBeInitializedStatically {
+        pushIfNotVisited(f)
+      }
+    }
+  }
+
   mutating func add(calleesOf function: Function) {
     for inst in function.instructions {
       switch inst {

SwiftCompilerSources/Sources/SIL/GlobalVariable.swift

@@ -53,6 +53,10 @@ final public class GlobalVariable : CustomStringConvertible, HasShortDescription
     return bridged.canBeInitializedStatically()
   }
 
+  public var mustBeInitializedStatically: Bool {
+    return bridged.mustBeInitializedStatically()
+  }
+
   public static func ==(lhs: GlobalVariable, rhs: GlobalVariable) -> Bool {
     lhs === rhs
   }

include/swift/AST/DiagnosticsSIL.def

@@ -342,6 +342,8 @@ ERROR(performance_callee_unavailable,none,
       "called function is not available in this module and can have unpredictable performance", ())
 ERROR(bad_attr_on_non_const_global,none,
       "global variable must be a compile-time constant to use %0 attribute", (StringRef))
+ERROR(global_must_be_compile_time_const,none,
+      "global variable must be a compile-time constant", ())
 NOTE(performance_called_from,none,
       "called from here", ())
 

include/swift/SIL/SILBridging.h

@@ -369,6 +369,8 @@ struct BridgedGlobalVar {
   inline OptionalBridgedInstruction getStaticInitializerValue() const;
 
   bool canBeInitializedStatically() const;
+  
+  bool mustBeInitializedStatically() const;
 };
 
 struct OptionalBridgedGlobalVar {

include/swift/SIL/SILGlobalVariable.h

@@ -173,6 +173,8 @@ class SILGlobalVariable
   /// static initializer.
   SILInstruction *getStaticInitializerValue();
 
+  bool mustBeInitializedStatically() const;
+
   /// Returns true if the global is a statically initialized heap object.
   bool isInitializedObject() {
     return dyn_cast_or_null<ObjectInst>(getStaticInitializerValue()) != nullptr;

lib/SIL/IR/SILGlobalVariable.cpp

@@ -81,6 +81,17 @@ SILInstruction *SILGlobalVariable::getStaticInitializerValue() {
   return &StaticInitializerBlock.back();
 }
 
+bool SILGlobalVariable::mustBeInitializedStatically() const {
+  if (getSectionAttr())
+    return true;
+
+  auto *decl = getDecl();
+  if (decl && isDefinition() && decl->getAttrs().hasAttribute<SILGenNameAttr>())
+    return true;
+
+  return false;
+}
+
 /// Return whether this variable corresponds to a Clang node.
 bool SILGlobalVariable::hasClangNode() const {
   return (VDecl ? VDecl->hasClangNode() : false);

lib/SIL/Utils/SILBridging.cpp

@@ -208,6 +208,11 @@ bool BridgedGlobalVar::canBeInitializedStatically() const {
   return tl.isLoadable();
 }
 
+bool BridgedGlobalVar::mustBeInitializedStatically() const {
+  SILGlobalVariable *global = getGlobal();
+  return global->mustBeInitializedStatically();
+}
+
 //===----------------------------------------------------------------------===//
 //                            SILVTable
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/Mandatory/PerformanceDiagnostics.cpp

@@ -483,18 +483,20 @@ class PerformanceDiagnosticsPass : public SILModuleTransform {
 
     // Check that @_section, @_silgen_name is only on constant globals
     for (SILGlobalVariable &g : module->getSILGlobals()) {
-      if (!g.getStaticInitializerValue()) {
+      if (!g.getStaticInitializerValue() && g.mustBeInitializedStatically()) {
+        auto *decl = g.getDecl();
         if (g.getSectionAttr()) {
           module->getASTContext().Diags.diagnose(
             g.getDecl()->getLoc(), diag::bad_attr_on_non_const_global,
             "@_section");
-        }
-
-        auto *decl = g.getDecl();
-        if (decl && g.isDefinition() && decl->getAttrs().hasAttribute<SILGenNameAttr>()) {
+        } else if (decl && g.isDefinition() &&
+                   decl->getAttrs().hasAttribute<SILGenNameAttr>()) {
           module->getASTContext().Diags.diagnose(
             g.getDecl()->getLoc(), diag::bad_attr_on_non_const_global,
             "@_silgen_name");
+        } else {
+          module->getASTContext().Diags.diagnose(
+            g.getDecl()->getLoc(), diag::global_must_be_compile_time_const);
         }
       }
     }