Merge pull request #81441 from eeckstein/vector_base_addr

SIL: introduce the `vector_base_addr` instruction and use it in `InlineArray`

Author: eeckstein

SwiftCompilerSources/Sources/AST/Type.swift

@@ -61,6 +61,8 @@ public struct Type: TypeProperties, CustomStringConvertible, NoReflectionChildre
 
   public var builtinVectorElementType: Type { Type(bridged: bridged.getBuiltinVectorElementType()) }
 
+  public var builtinFixedArrayElementType: Type { Type(bridged: bridged.getBuiltinFixedArrayElementType()) }
+
   public func subst(with substitutionMap: SubstitutionMap) -> Type {
     return Type(bridged: bridged.subst(substitutionMap.bridged))
   }
@@ -81,6 +83,8 @@ public struct CanonicalType: TypeProperties, CustomStringConvertible, NoReflecti
 
   public var builtinVectorElementType: CanonicalType { rawType.builtinVectorElementType.canonical }
 
+  public var builtinFixedArrayElementType: CanonicalType { rawType.builtinFixedArrayElementType.canonical }
+
   public func subst(with substitutionMap: SubstitutionMap) -> CanonicalType {
     return rawType.subst(with: substitutionMap).canonical
   }
@@ -106,6 +110,7 @@ extension TypeProperties {
 
   public var isBuiltinFloat: Bool { rawType.bridged.isBuiltinFloat() }
   public var isBuiltinVector: Bool { rawType.bridged.isBuiltinVector() }
+  public var isBuiltinFixedArray: Bool { rawType.bridged.isBuiltinFixedArray() }
 
   public var isClass: Bool {
     if let nominal = nominal, nominal is ClassDecl {

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/InitializeStaticGlobals.swift

@@ -175,7 +175,7 @@ private func lowerInlineArray(array: InlineArray, _ context: FunctionPassContext
 ///
 private func getInlineArrayInfo(of allocStack: AllocStackInst) -> InlineArray? {
   var arrayLoad: LoadInst? = nil
-  var elementStorage: UncheckedAddrCastInst? = nil
+  var elementStorage: VectorBaseAddrInst? = nil
 
   for use in allocStack.uses {
     switch use.instruction {
@@ -188,11 +188,11 @@ private func getInlineArrayInfo(of allocStack: AllocStackInst) -> InlineArray? {
       arrayLoad = load
     case is DeallocStackInst:
       break
-    case let addrCastToElement as UncheckedAddrCastInst:
+    case let baseAddr as VectorBaseAddrInst:
       if elementStorage != nil {
         return nil
       }
-      elementStorage = addrCastToElement
+      elementStorage = baseAddr
     default:
       return nil
     }

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/CMakeLists.txt

@@ -39,6 +39,7 @@ swift_compiler_sources(Optimizer
   SimplifySwitchEnum.swift
   SimplifyTuple.swift
   SimplifyTupleExtract.swift
+  SimplifyUncheckedAddrCast.swift
   SimplifyUncheckedEnumData.swift
   SimplifyValueToBridgeObject.swift
   SimplifyWitnessMethod.swift)

SwiftCompilerSources/Sources/Optimizer/InstructionSimplification/SimplifyUncheckedAddrCast.swift

@@ -0,0 +1,81 @@
+//===--- SimplifyUncheckedAddrCast.swift ----------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2025 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SIL
+
+extension UncheckedAddrCastInst : OnoneSimplifiable, SILCombineSimplifiable {
+
+  func simplify(_ context: SimplifyContext) {
+    // ```
+    //   %1 = unchecked_addr_cast %0 : $*T to $*T
+    // ```
+    // ->
+    //   replace %1 with %0
+    //
+    if optimizeSameTypeCast(context) {
+      return
+    }
+
+    // ```
+    //   %1 = unchecked_addr_cast %0 : $*U to $*V
+    //   %2 = unchecked_addr_cast %1 : $*V to $*T
+    // ```
+    // ->
+    // ```
+    //   %2 = unchecked_addr_cast %0: $*U to $*T
+    // ```
+    if optimizeDoubleCast(context) {
+      return
+    }
+
+    // ```
+    //   %1 = unchecked_addr_cast %0 : $*Builtin.FixedArray<N, Element> to $*Element
+    // ```
+    // ->
+    // ```
+    //   %1 = vector_base_addr %0 : $*Builtin.FixedArray<N, Element>
+    // ```
+    _ = optimizeVectorBaseCast(context)
+  }
+}
+
+private extension UncheckedAddrCastInst {
+  func optimizeSameTypeCast(_ context: SimplifyContext) -> Bool {
+    if fromAddress.type == type {
+      self.replace(with: fromAddress, context)
+      return true
+    }
+    return false
+  }
+
+  func optimizeDoubleCast(_ context: SimplifyContext) -> Bool {
+    if let firstCast = fromAddress as? UncheckedAddrCastInst {
+      let builder = Builder(before: self, context)
+      let newCast = builder.createUncheckedAddrCast(from: firstCast.fromAddress, to: type)
+      self.replace(with: newCast, context)
+      return true
+    }
+    return false
+  }
+
+  func optimizeVectorBaseCast(_ context: SimplifyContext) -> Bool {
+    if fromAddress.type.isBuiltinFixedArray,
+       fromAddress.type.builtinFixedArrayElementType(in: parentFunction, maximallyAbstracted: true).addressType == type
+    {
+      let builder = Builder(before: self, context)
+      let vectorBase = builder.createVectorBaseAddr(vector: fromAddress)
+      self.replace(with: vectorBase, context)
+      return true
+    }
+    return false
+  }
+}

SwiftCompilerSources/Sources/Optimizer/PassManager/PassRegistration.swift

@@ -125,6 +125,7 @@ private func registerSwiftPasses() {
   registerForSILCombine(PointerToAddressInst.self,  { run(PointerToAddressInst.self, $0) })
   registerForSILCombine(UncheckedEnumDataInst.self, { run(UncheckedEnumDataInst.self, $0) })
   registerForSILCombine(WitnessMethodInst.self,     { run(WitnessMethodInst.self, $0) })
+  registerForSILCombine(UncheckedAddrCastInst.self, { run(UncheckedAddrCastInst.self, $0) })
   registerForSILCombine(UnconditionalCheckedCastInst.self, { run(UnconditionalCheckedCastInst.self, $0) })
   registerForSILCombine(AllocStackInst.self,        { run(AllocStackInst.self, $0) })
   registerForSILCombine(ApplyInst.self,             { run(ApplyInst.self, $0) })

SwiftCompilerSources/Sources/Optimizer/Utilities/OptUtils.swift

@@ -387,7 +387,7 @@ extension Instruction {
     case let bi as BuiltinInst:
       switch bi.id {
       case .ZeroInitializer:
-        let type = bi.type.isBuiltinVector ? bi.type.builtinVectorElementType : bi.type
+        let type = bi.type.isBuiltinVector ? bi.type.builtinVectorElementType(in: parentFunction) : bi.type
         return type.isBuiltinInteger || type.isBuiltinFloat
       case .PtrToInt:
         return bi.operands[0].value is StringLiteralInst

SwiftCompilerSources/Sources/Optimizer/Utilities/Verifier.swift

@@ -148,6 +148,16 @@ extension LoadBorrowInst : VerifiableInstruction {
   }
 }
 
+extension VectorBaseAddrInst : VerifiableInstruction {
+  func verify(_ context: FunctionPassContext) {
+    require(vector.type.isBuiltinFixedArray,
+            "vector operand of vector_element_addr must be a Builtin.FixedArray")
+    require(type == vector.type.builtinFixedArrayElementType(in: parentFunction,
+                                                             maximallyAbstracted: true).addressType,
+            "result of vector_element_addr has wrong type")
+  }
+}
+
 // Used to check if any instruction is mutating the memory location within the liverange of a `load_borrow`.
 // Note that it is not checking if an instruction _may_ mutate the memory, but it's checking if any instruction
 // _definitely_ will mutate the memory.

SwiftCompilerSources/Sources/SIL/Builder.swift

@@ -487,6 +487,10 @@ public struct Builder {
     return notifyNew(vectorInst.getAs(VectorInst.self))
   }
 
+  public func createVectorBaseAddr(vector: Value) -> VectorBaseAddrInst {
+    return notifyNew(bridged.createVectorBaseAddr(vector.bridged).getAs(VectorBaseAddrInst.self))
+  }
+
   public func createGlobalAddr(global: GlobalVariable, dependencyToken: Value?) -> GlobalAddrInst {
     return notifyNew(bridged.createGlobalAddr(global.bridged, dependencyToken.bridged).getAs(GlobalAddrInst.self))
   }

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1350,6 +1350,10 @@ final public class ObjectInst : SingleValueInstruction {
 final public class VectorInst : SingleValueInstruction {
 }
 
+final public class VectorBaseAddrInst : SingleValueInstruction, UnaryInstruction {
+  public var vector: Value { operand.value }
+}
+
 final public class DifferentiableFunctionInst: SingleValueInstruction {}
 
 final public class LinearFunctionInst: SingleValueInstruction {}

SwiftCompilerSources/Sources/SIL/Registration.swift

@@ -118,6 +118,7 @@ public func registerSILClasses() {
   register(MoveOnlyWrapperToCopyableAddrInst.self)
   register(ObjectInst.self)
   register(VectorInst.self)
+  register(VectorBaseAddrInst.self)
   register(TuplePackExtractInst.self)
   register(TuplePackElementAddrInst.self)
   register(PackElementGetInst.self)