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

kubamracek · kubamracek · commit 145f12f6a30e · 2023-07-08T19:26:59.000-07:00
Before this change, if a global variable is required to be statically initialized (e.g. due to @_section attribute), we don't allow its type to be a struct, only a scalar type works. This change improves on that by teaching MandatoryPerformanceOptimizations pass to inline struct initializer calls into initializer of globals, as long as they are simple enough so that we can be sure that we don't trigger recursive/infinite inlining.
diff --git a/SwiftCompilerSources/Sources/Optimizer/FunctionPasses/SimplificationPasses.swift b/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,
diff --git a/SwiftCompilerSources/Sources/Optimizer/ModulePasses/MandatoryPerformanceOptimizations.swift b/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)
 }
@@ -48,33 +50,37 @@ private func optimizeFunctionsTopDown(using worklist: inout FunctionWorklist,
 }
 
 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: [SmallProjectionPath:Set<Function>] = [:]
+  
+  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 [SmallProjectionPath:Set<Function>],
+                                   _ context: FunctionPassContext, _ simplifyCtxt: SimplifyContext) {
   if simplifyCtxt.tryDevirtualize(apply: apply, isMandatory: true) != nil {
     return
   }
@@ -88,7 +94,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 +116,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 [SmallProjectionPath:Set<Function>]) -> Bool {
   if callee.isTransparent {
     return true
   }
@@ -123,9 +129,104 @@ 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[projectionPath, default: Set()].contains(callee) {
+    alreadyInlinedFunctions[projectionPath, default: Set()].insert(callee)
+    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.singleUse 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 is EnumInst:
+        path = path.push(.enumCase, index: use.index)
+        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:
+        break
+      }
+
+      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 +247,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 {
diff --git a/SwiftCompilerSources/Sources/SIL/GlobalVariable.swift b/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
   }
diff --git a/include/swift/AST/DiagnosticsSIL.def b/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", ())
 
diff --git a/include/swift/SIL/SILBridging.h b/include/swift/SIL/SILBridging.h
@@ -369,6 +369,8 @@ struct BridgedGlobalVar {
   inline OptionalBridgedInstruction getStaticInitializerValue() const;
 
   bool canBeInitializedStatically() const;
+  
+  bool mustBeInitializedStatically() const;
 };
 
 struct OptionalBridgedGlobalVar {
diff --git a/include/swift/SIL/SILGlobalVariable.h b/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;
diff --git a/lib/SIL/IR/SILGlobalVariable.cpp b/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);
diff --git a/lib/SIL/Utils/SILBridging.cpp b/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
 //===----------------------------------------------------------------------===//
diff --git a/lib/SILOptimizer/Mandatory/PerformanceDiagnostics.cpp b/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);
         }
       }
     }
diff --git a/test/IRGen/section_non_const.swift b/test/IRGen/section_non_const.swift
diff --git a/test/IRGen/section_structs.swift b/test/IRGen/section_structs.swift

Original file line number	Diff line number	Diff line change
`@@ -53,6 +53,10 @@ final public class GlobalVariable : CustomStringConvertible, HasShortDescription`
`53`	`53`	`return bridged.canBeInitializedStatically()`
`54`	`54`	`}`
`55`	`55`
	`56`	`+ public var mustBeInitializedStatically: Bool {`
	`57`	`+ return bridged.mustBeInitializedStatically()`
	`58`	`+ }`
	`59`	`+`
`56`	`60`	`public static func ==(lhs: GlobalVariable, rhs: GlobalVariable) -> Bool {`
`57`	`61`	`lhs === rhs`
`58`	`62`	`}`