Merge pull request #79298 from eeckstein/static-init-inline-arrays

Support statically initializing globals which are or contain inline arrays

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -58,19 +58,21 @@ let initializeStaticGlobalsPass = FunctionPass(name: "initialize-static-globals"
   // Merge such individual stores to a single store of the whole struct.
   mergeStores(in: function, context)
 
-  // The initializer must not contain a `global_value` because `global_value` needs to
-  // initialize the class metadata at runtime.
-  guard let (allocInst, storeToGlobal) = getGlobalInitialization(of: function,
-                                                                 forStaticInitializer: true,
-                                                                 context) else
-  {
+  guard let (allocInst, storeToGlobal, inlineArrays) = getGlobalInitializerInfo(of: function, context) else {
     return
   }
 
   if !allocInst.global.canBeInitializedStatically {
     return
   }
 
+  /// Replace inline arrays, which are allocated in stack locations with `vector` instructions.
+  /// Note that `vector` instructions are only allowed in global initializers. Therefore it's important
+  /// that the code in this global initializer is eventually completely removed after copying it to the global.
+  for array in inlineArrays {
+    lowerInlineArray(array: array, context)
+  }
+    
   var cloner = StaticInitCloner(cloneTo: allocInst.global, context)
   defer { cloner.deinitialize() }
 
@@ -87,6 +89,186 @@ let initializeStaticGlobalsPass = FunctionPass(name: "initialize-static-globals"
   context.removeTriviallyDeadInstructionsIgnoringDebugUses(in: function)
 }
 
+/// Gets all info about a global initializer function if it can be converted to a statically initialized global.
+private func getGlobalInitializerInfo(
+  of function: Function,
+  _ context: FunctionPassContext
+) -> (allocInst: AllocGlobalInst, storeToGlobal: StoreInst, inlineArrays: [InlineArray])? {
+
+  var arrayInitInstructions = InstructionSet(context)
+  defer { arrayInitInstructions.deinitialize() }
+  
+  var inlineArrays = [InlineArray]()
+  
+  guard let (allocInst, storeToGlobal) = getGlobalInitialization(of: function, context,
+    handleUnknownInstruction: { inst in
+      if let asi = inst as? AllocStackInst {
+        if let array = getInlineArrayInfo(of: asi) {
+          inlineArrays.append(array)
+          arrayInitInstructions.insertAllAddressUses(of: asi)
+          return true
+        }
+        return false
+      }
+      // Accept all instructions which are part of inline array initialization, because we'll remove them anyway.
+      return arrayInitInstructions.contains(inst)
+    })
+  else {
+    return nil
+  }
+
+  return (allocInst, storeToGlobal, inlineArrays)   
+}
+
+/// Represents an inline array which is initialized by a literal.
+private struct InlineArray {
+  let elementType: Type
+  
+  /// In case the `elementType` is a tuple, the element values are flattened,
+  /// i.e. `elements` contains elementcount * tupleelements values.
+  let elements: [Value]
+  
+  /// The final load instruction which loads the initialized array from a temporary stack location.
+  let finalArrayLoad: LoadInst
+  
+  /// The stack location which contains the initialized array.
+  var stackLoocation: AllocStackInst { finalArrayLoad.address as! AllocStackInst }
+}
+
+/// Replaces an initialized inline array (which is allocated in a temporary stack location) with a
+/// `vector` instruction.
+/// The stack location of the array is removed.
+private func lowerInlineArray(array: InlineArray, _ context: FunctionPassContext) {
+  let vector: VectorInst
+  let builder = Builder(after: array.finalArrayLoad, context)
+  if array.elementType.isTuple {
+    let numTupleElements = array.elementType.tupleElements.count
+    assert(array.elements.count % numTupleElements == 0)
+    var tuples: [TupleInst] = []
+    for tupleIdx in 0..<(array.elements.count / numTupleElements) {
+      let range = (tupleIdx * numTupleElements) ..< ((tupleIdx + 1) * numTupleElements) 
+      let tuple = builder.createTuple(type: array.elementType, elements: Array(array.elements[range]))
+      tuples.append(tuple)
+    }
+    vector = builder.createVector(type: array.elementType, arguments: tuples)
+  } else {
+    vector = builder.createVector(type: array.elementType, arguments: array.elements)      
+  }
+  array.finalArrayLoad.uses.replaceAll(with: vector, context)
+  context.erase(instructionIncludingAllUsers: array.stackLoocation)
+}
+
+/// An alloc_stack could be a temporary object which holds an initialized inline-array literal.
+/// It looks like:
+///
+///     %1 = alloc_stack $InlineArray<Count, ElementType>
+///     %2 = unchecked_addr_cast %1 to $*ElementType      // the elementStorage
+///     store %firstElement to [trivial] %2
+///     %4 = integer_literal $Builtin.Word, 1
+///     %5 = index_addr %2, %4
+///     store %secondElement to [trivial] %5
+///     ...
+///     %10 = load [trivial] %1                   // the final arrayLoad
+///     dealloc_stack %1
+///
+/// Returns nil if `allocStack` is not a properly initialized inline array.
+///
+private func getInlineArrayInfo(of allocStack: AllocStackInst) -> InlineArray? {
+  var arrayLoad: LoadInst? = nil
+  var elementStorage: UncheckedAddrCastInst? = nil
+
+  for use in allocStack.uses {
+    switch use.instruction {
+    case let load as LoadInst:
+      if arrayLoad != nil {
+        return nil
+      }
+      // It's guaranteed that the array load is located after all element stores.
+      // Otherwise it would load uninitialized memory.
+      arrayLoad = load
+    case is DeallocStackInst:
+      break
+    case let addrCastToElement as UncheckedAddrCastInst:
+      if elementStorage != nil {
+        return nil
+      }
+      elementStorage = addrCastToElement
+    default:
+      return nil
+    }
+  }
+  guard let arrayLoad, let elementStorage else {
+    return nil
+  } 
+  
+  var stores = Array<StoreInst?>()
+  if !findArrayElementStores(toElementAddress: elementStorage, elementIndex: 0, stores: &stores) {
+    return nil
+  }
+  if stores.isEmpty {
+    // We cannot create an empty `vector` instruction, therefore we don't support empty inline arrays.
+    return nil
+  }
+  // Usually there must be a store for each element. Otherwise the `arrayLoad` would load uninitialized memory.
+  // We still check this to not crash in some weird corner cases, like the element type is an empty tuple.
+  if stores.contains(nil) {
+    return nil
+  }
+
+  return InlineArray(elementType: elementStorage.type.objectType,
+                     elements: stores.map { $0!.source },
+                     finalArrayLoad: arrayLoad)
+}
+
+/// Recursively traverses all uses of `elementAddr` and finds all stores to an inline array storage.
+/// The element store instructions are put into `stores` - one store for each element.
+/// In case the element type is a tuple, the tuples are flattened. See `InlineArray.elements`.  
+private func findArrayElementStores(
+  toElementAddress elementAddr: Value,
+  elementIndex: Int,
+  stores: inout [StoreInst?]
+) -> Bool {
+  for use in elementAddr.uses {
+    switch use.instruction {
+    case let indexAddr as IndexAddrInst:
+      guard let indexLiteral = indexAddr.index as? IntegerLiteralInst,
+            let tailIdx = indexLiteral.value else
+      {
+        return false
+      }
+      if !findArrayElementStores(toElementAddress: indexAddr, elementIndex: elementIndex + tailIdx, stores: &stores) {
+        return false
+      }
+    case let tea as TupleElementAddrInst:
+      // The array elements are tuples. There is a separate store for each tuple element.
+      let numTupleElements = tea.tuple.type.tupleElements.count
+      let tupleIdx = tea.fieldIndex
+      if !findArrayElementStores(toElementAddress: tea,
+                                 elementIndex: elementIndex * numTupleElements + tupleIdx,
+                                 stores: &stores) {
+        return false
+      }
+    case let store as StoreInst:
+      if store.source.type.isTuple {
+        // This kind of SIL is never generated because tuples are stored with separated stores to tuple_element_addr.
+        // Just to be on the safe side..
+        return false
+      }
+      if elementIndex >= stores.count {
+        stores += Array(repeating: nil, count: elementIndex - stores.count + 1)
+      }
+      if stores[elementIndex] != nil {
+        // An element is stored twice.
+        return false
+      }
+      stores[elementIndex] = store
+    default:
+      return false
+    }
+  }
+  return true
+}
+
 /// Merges stores to individual struct fields to a single store of the whole struct.
 ///
 ///   store %element1 to %element1Addr
@@ -172,3 +354,15 @@ private func merge(elementStores: [StoreInst], lastStore: StoreInst, _ context:
     }
   }
 }
+
+private extension InstructionSet {
+  mutating func insertAllAddressUses(of value: Value) {
+    for use in value.uses {
+      if insert(use.instruction) {
+        for result in use.instruction.results where result.type.isAddress {
+          insertAllAddressUses(of: result)
+        }
+      }
+    }
+  }
+}

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyLoad.swift

@@ -272,7 +272,10 @@ private func getInitializerFromInitFunction(of globalAddr: GlobalAddrInst, _ con
   }
   let initFn = initFnRef.referencedFunction
   context.notifyDependency(onBodyOf: initFn)
-  guard let (_, storeToGlobal) = getGlobalInitialization(of: initFn, forStaticInitializer: false, context) else {
+  guard let (_, storeToGlobal) = getGlobalInitialization(of: initFn, context, handleUnknownInstruction: {
+    // Accept `global_value` because the class header can be initialized at runtime by the `global_value` instruction.
+    return $0 is GlobalValueInst
+  }) else {
     return nil
   }
   return storeToGlobal.source
@@ -305,10 +308,6 @@ private func transitivelyErase(load: LoadInst, _ context: SimplifyContext) {
 
 private extension Value {
   func canBeCopied(into function: Function, _ context: SimplifyContext) -> Bool {
-    if !function.isAnySerialized {
-      return true
-    }
-
     // Can't use `ValueSet` because the this value is inside a global initializer and
     // not inside a function.
     var worklist = Stack<Value>(context)
@@ -320,8 +319,13 @@ private extension Value {
     handled.insert(ObjectIdentifier(self))
 
     while let value = worklist.pop() {
+      if value is VectorInst {
+        return false
+      }
       if let fri = value as? FunctionRefInst {
-        if !fri.referencedFunction.hasValidLinkageForFragileRef(function.serializedKind) {
+        if function.isAnySerialized, 
+           !fri.referencedFunction.hasValidLinkageForFragileRef(function.serializedKind)
+        {
           return false
         }
       }

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -791,10 +791,13 @@ extension InstructionRange {
 ///   %i = some_const_initializer_insts
 ///   store %i to %a
 /// ```
+/// 
+/// For all other instructions `handleUnknownInstruction` is called and such an instruction
+/// is accepted if `handleUnknownInstruction` returns true.
 func getGlobalInitialization(
   of function: Function,
-  forStaticInitializer: Bool,
-  _ context: some Context
+  _ context: some Context,
+  handleUnknownInstruction: (Instruction) -> Bool
 ) -> (allocInst: AllocGlobalInst, storeToGlobal: StoreInst)? {
   guard let block = function.blocks.singleElement else {
     return nil
@@ -811,34 +814,36 @@ func getGlobalInitialization(
          is DebugStepInst,
          is BeginAccessInst,
          is EndAccessInst:
-      break
+      continue
     case let agi as AllocGlobalInst:
-      if allocInst != nil {
-        return nil
+      if allocInst == nil {
+        allocInst = agi
+        continue
       }
-      allocInst = agi
     case let ga as GlobalAddrInst:
       if let agi = allocInst, agi.global == ga.global {
         globalAddr = ga
       }
+      continue
     case let si as StoreInst:
-      if store != nil {
-        return nil
-      }
-      guard let ga = globalAddr else {
-        return nil
-      }
-      if si.destination != ga {
-        return nil
+      if store == nil,
+         let ga = globalAddr,
+         si.destination == ga
+      {
+        store = si
+        continue
       }
-      store = si
-    case is GlobalValueInst where !forStaticInitializer:
-      break
+    // Note that the initializer must not contain a `global_value` because `global_value` needs to
+    // initialize the class metadata at runtime.
     default:
-      if !inst.isValidInStaticInitializerOfGlobal(context) {
-        return nil
+      if inst.isValidInStaticInitializerOfGlobal(context) {
+        continue
       }
     }
+    if handleUnknownInstruction(inst) {
+      continue
+    }
+    return nil
   }
   if let store = store {
     return (allocInst: allocInst!, storeToGlobal: store)

include/swift/SIL/SILBuilder.h

@@ -3115,12 +3115,16 @@ class SILBuilder {
     C.notifyInserted(TheInst);
 
 #ifndef NDEBUG
-    // If we are inserting into a specific function (rather than a block for a
-    // global_addr), verify that our instruction/the associated location are in
-    // sync. We don't care if an instruction is used in global_addr.
-    if (F)
-      TheInst->verifyDebugInfo();
-    TheInst->verifyOperandOwnership(&C.silConv);
+    // A vector instruction can only be in a global initializer. Therefore there
+    // is no point in verifying debug info or ownership.
+    if (!isa<VectorInst>(TheInst)) {
+      // If we are inserting into a specific function (rather than a block for a
+      // global_addr), verify that our instruction/the associated location are in
+      // sync. We don't care if an instruction is used in global_addr.
+      if (F)
+        TheInst->verifyDebugInfo();
+      TheInst->verifyOperandOwnership(&C.silConv);
+    }
 #endif
   }
 

lib/IRGen/GenConstant.cpp

@@ -433,6 +433,24 @@ Explosion irgen::emitConstantValue(IRGenModule &IGM, SILValue operand,
   } else if (auto *atp = dyn_cast<AddressToPointerInst>(operand)) {
     auto *val = emitConstantValue(IGM, atp->getOperand()).claimNextConstant();
     return val;
+  } else if (auto *vector = dyn_cast<VectorInst>(operand)) {
+    if (flatten) {
+      Explosion out;
+      for (SILValue element : vector->getElements()) {
+        Explosion e = emitConstantValue(IGM, element, flatten);
+        out.add(e.claimAll());
+      }
+      return out;
+    }
+    llvm::SmallVector<llvm::Constant *, 8> elementValues;
+    for (SILValue element : vector->getElements()) {
+      auto &ti = cast<FixedTypeInfo>(IGM.getTypeInfo(element->getType()));
+      Size paddingBytes = ti.getFixedStride() - ti.getFixedSize();
+      Explosion e = emitConstantValue(IGM, element, flatten);
+      elementValues.push_back(IGM.getConstantValue(std::move(e), paddingBytes.getValue()));
+    }
+    auto *arrTy = llvm::ArrayType::get(elementValues[0]->getType(), elementValues.size());
+    return llvm::ConstantArray::get(arrTy, elementValues);
   } else {
     llvm_unreachable("Unsupported SILInstruction in static initializer!");
   }