Merge pull request #83999 from eeckstein/access-enforcement-improvements

Optimizer: some small fixes and improvements for exclusivity checking

Author: eeckstein

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/AllocBoxToStack.swift

@@ -54,7 +54,7 @@ import SIL
 let allocBoxToStack = FunctionPass(name: "allocbox-to-stack") {
   (function: Function, context: FunctionPassContext) in
 
-  _ = tryConvertBoxesToStack(in: function, context)
+  _ = tryConvertBoxesToStack(in: function, isMandatory: false, context)
 }
 
 /// The "mandatory" version of the pass, which runs in the mandatory pipeline.
@@ -72,7 +72,7 @@ let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
 
   while let function = worklist.pop() {
     moduleContext.transform(function: function) { context in
-      let specFns = tryConvertBoxesToStack(in: function, context)
+      let specFns = tryConvertBoxesToStack(in: function, isMandatory: true, context)
       worklist.pushIfNotVisited(contentsOf: specFns.specializedFunctions)
       originalsOfSpecializedFunctions.pushIfNotVisited(contentsOf: specFns.originalFunctions)
     }
@@ -86,9 +86,11 @@ let mandatoryAllocBoxToStack = ModulePass(name: "mandatory-allocbox-to-stack") {
 /// Converts all non-escaping `alloc_box` to `alloc_stack` and specializes called functions if a
 /// box is passed to a function.
 /// Returns the list of original functions for which a specialization has been created.
-private func tryConvertBoxesToStack(in function: Function, _ context: FunctionPassContext) -> FunctionSpecializations {
+private func tryConvertBoxesToStack(in function: Function, isMandatory: Bool,
+                                    _ context: FunctionPassContext
+) -> FunctionSpecializations {
   var promotableBoxes = Array<(AllocBoxInst, Flags)>()
-  var functionsToSpecialize = FunctionSpecializations()
+  var functionsToSpecialize = FunctionSpecializations(isMandatory: isMandatory)
 
   findPromotableBoxes(in: function, &promotableBoxes, &functionsToSpecialize)
 
@@ -189,6 +191,9 @@ private struct FunctionSpecializations {
   private var promotableArguments = CrossFunctionValueWorklist()
   private var originals = FunctionWorklist()
   private var originalToSpecialized = Dictionary<Function, Function>()
+  private let isMandatory: Bool
+
+  init(isMandatory: Bool) { self.isMandatory = isMandatory }
 
   var originalFunctions: [Function] { originals.functions }
   var specializedFunctions: [Function] { originals.functions.lazy.map { originalToSpecialized[$0]! } }
@@ -221,6 +226,12 @@ private struct FunctionSpecializations {
         context.erase(instruction: user)
       case let projectBox as ProjectBoxInst:
         assert(projectBox.fieldIndex == 0, "only single-field boxes are handled")
+        if isMandatory {
+          // Once we have promoted the box to stack, access violations can be detected statically by the
+          // DiagnoseStaticExclusivity pass (which runs after MandatoryAllocBoxToStack).
+          // Therefore we can convert dynamic accesses to static accesses.
+          makeAccessesStatic(of: projectBox, context)
+        }
         projectBox.replace(with: stack, context)
       case is MarkUninitializedInst, is CopyValueInst, is BeginBorrowInst, is MoveValueInst:
         // First, replace the instruction with the original `box`, which adds more uses to `box`.
@@ -478,6 +489,14 @@ private func hoistMarkUnresolvedInsts(stackAddress: Value,
     .replaceAll(with: mu, context)
 }
 
+private func makeAccessesStatic(of address: Value, _ context: FunctionPassContext) {
+  for beginAccess in address.uses.users(ofType: BeginAccessInst.self) {
+    if beginAccess.enforcement == .dynamic {
+      beginAccess.set(enforcement: .static, context: context)
+    }
+  }
+}
+
 private extension ApplySite {
   func getSpecializableCallee() -> Function? {
     if let callee = referencedFunction,

SwiftCompilerSources/Sources/Optimizer/FunctionPasses/TempLValueElimination.swift

@@ -113,6 +113,13 @@ private func tryEliminate(copy: CopyLikeInstruction, _ context: FunctionPassCont
     return
   }
 
+  // If exclusivity is checked for the alloc_stack we must not replace it with the copy-destination.
+  // If the copy-destination is also in an access-scope this would result in an exclusivity violation
+  // which was not there before.
+  guard allocStack.uses.users(ofType: BeginAccessInst.self).isEmpty else {
+    return
+  }
+
   var worklist = InstructionWorklist(context)
   defer { worklist.deinitialize() }
   worklist.pushIfNotVisited(firstUseOfAllocStack)

SwiftCompilerSources/Sources/SIL/Instruction.swift

@@ -1629,12 +1629,39 @@ final public class BeginAccessInst : SingleValueInstruction, UnaryInstruction {
   public var accessKind: AccessKind {
     AccessKind(rawValue: bridged.BeginAccessInst_getAccessKind())!
   }
-  public func set(accessKind: BeginAccessInst.AccessKind, context: some MutatingContext) {
+  public func set(accessKind: AccessKind, context: some MutatingContext) {
     context.notifyInstructionsChanged()
     bridged.BeginAccess_setAccessKind(accessKind.rawValue)
     context.notifyInstructionChanged(self)
   }
 
+  // The raw values must match SILAccessEnforcement.
+  public enum Enforcement: Int {
+    /// The access's enforcement has not yet been determined.
+    case unknown = 0
+
+    /// The access is statically known to not conflict with other accesses.
+    case `static` = 1
+
+    /// The access is not statically known to not conflict with anything and must be dynamically checked.
+    case dynamic = 2
+
+    /// The access is not statically known to not conflict with anything but dynamic checking should
+    /// be suppressed, leaving it undefined behavior.
+    case unsafe = 3
+
+    /// Access to pointers that are signed via pointer authentication.
+    case signed = 4
+  }
+  public var enforcement: Enforcement {
+    Enforcement(rawValue: bridged.BeginAccessInst_getEnforcement())!
+  }
+  public func set(enforcement: Enforcement, context: some MutatingContext) {
+    context.notifyInstructionsChanged()
+    bridged.BeginAccess_setEnforcement(enforcement.rawValue)
+    context.notifyInstructionChanged(self)
+  }
+
   public var isStatic: Bool { bridged.BeginAccessInst_isStatic() }
   public var isUnsafe: Bool { bridged.BeginAccessInst_isUnsafe() }
 

include/swift/SIL/SILBridging.h

@@ -824,6 +824,8 @@ struct BridgedInstruction {
   BRIDGED_INLINE bool BeginAccessInst_isStatic() const;
   BRIDGED_INLINE bool BeginAccessInst_isUnsafe() const;
   BRIDGED_INLINE void BeginAccess_setAccessKind(SwiftInt accessKind) const;
+  BRIDGED_INLINE SwiftInt BeginAccessInst_getEnforcement() const;
+  BRIDGED_INLINE void BeginAccess_setEnforcement(SwiftInt accessKind) const;
   BRIDGED_INLINE bool CopyAddrInst_isTakeOfSrc() const;
   BRIDGED_INLINE bool CopyAddrInst_isInitializationOfDest() const;
   BRIDGED_INLINE void CopyAddrInst_setIsTakeOfSrc(bool isTakeOfSrc) const;

include/swift/SIL/SILBridgingImpl.h

@@ -1494,6 +1494,14 @@ void BridgedInstruction::BeginAccess_setAccessKind(SwiftInt accessKind) const {
   getAs<swift::BeginAccessInst>()->setAccessKind((swift::SILAccessKind)accessKind);
 }
 
+SwiftInt BridgedInstruction::BeginAccessInst_getEnforcement() const {
+  return (SwiftInt)getAs<swift::BeginAccessInst>()->getEnforcement();
+}
+
+void BridgedInstruction::BeginAccess_setEnforcement(SwiftInt accessKind) const {
+  getAs<swift::BeginAccessInst>()->setEnforcement((swift::SILAccessEnforcement)accessKind);
+}
+
 bool BridgedInstruction::CopyAddrInst_isTakeOfSrc() const {
   return getAs<swift::CopyAddrInst>()->isTakeOfSrc();
 }

include/swift/SIL/SILInstruction.h

@@ -4885,6 +4885,7 @@ class EndBorrowInst
 };
 
 /// Different kinds of access.
+/// This enum must stay in sync with `BeginAccessInst.AccessKind` in SwiftCompilerSources.
 enum class SILAccessKind : uint8_t {
   /// An access which takes uninitialized memory and initializes it.
   Init,
@@ -4907,6 +4908,7 @@ enum { NumSILAccessKindBits = 2 };
 StringRef getSILAccessKindName(SILAccessKind kind);
 
 /// Different kinds of exclusivity enforcement for accesses.
+/// This enum must stay in sync with `BeginAccessInst.Enforcement` in SwiftCompilerSources.
 enum class SILAccessEnforcement : uint8_t {
   /// The access's enforcement has not yet been determined.
   Unknown,

lib/SILOptimizer/Mandatory/RawSILInstLowering.cpp

@@ -290,13 +290,20 @@ lowerAssignOrInitInstruction(SILBuilderWithScope &b,
       auto isRefSelf = selfValue->getType().getASTType()->mayHaveSuperclass();
 
       SILValue selfRef;
+      bool needInsertEndAccess = false;
       if (isRefSelf) {
         selfRef = b.emitBeginBorrowOperation(loc, selfValue);
+      } else if (isa<BeginAccessInst>(selfValue)) {
+        // Don't insert an access scope if there is already one. This avoids
+        // inserting a dynamic access check when the parent is static (and therefore
+        // can be statically enforced).
+        selfRef = selfValue;
       } else {
         selfRef = b.createBeginAccess(loc, selfValue, SILAccessKind::Modify,
                                       SILAccessEnforcement::Dynamic,
                                       /*noNestedConflict=*/false,
                                       /*fromBuiltin=*/false);
+        needInsertEndAccess = true;
       }
 
       auto emitFieldReference = [&](VarDecl *field,
@@ -341,7 +348,7 @@ lowerAssignOrInitInstruction(SILBuilderWithScope &b,
       if (isRefSelf) {
         if (selfRef != selfValue)
           b.emitEndBorrowOperation(loc, selfRef);
-      } else {
+      } else if (needInsertEndAccess) {
         b.createEndAccess(loc, selfRef, /*aborted=*/false);
       }
 

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -117,7 +117,6 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
   P.startPipeline("Mandatory Diagnostic Passes + Enabling Optimization Passes");
   P.addDiagnoseInvalidEscapingCaptures();
   P.addReferenceBindingTransform();
-  P.addDiagnoseStaticExclusivity();
   P.addNestedSemanticFunctionCheck();
   P.addCapturePromotion();
 
@@ -130,6 +129,10 @@ static void addMandatoryDiagnosticOptPipeline(SILPassPipelinePlan &P) {
 #else
   P.addLegacyAllocBoxToStack();
 #endif
+  // Needs to run after MandatoryAllocBoxToStack, because MandatoryAllocBoxToStack
+  // can convert dynamic accesses to static accesses.
+  P.addDiagnoseStaticExclusivity();
+
   P.addNoReturnFolding();
   P.addBooleanLiteralFolding();
   addDefiniteInitialization(P);

test/SILGen/moveonly_escaping_closure.swift

@@ -143,7 +143,7 @@ func testLocalLetClosureCaptureVar() {
         consumeVal(x) // expected-note {{consumed here}}
         // expected-note @-1 {{consumed again here}}
         borrowConsumeVal(x, x)
-        // expected-error @-1 {{overlapping accesses, but deinitialization requires exclusive access}}
+        // expected-error @-1 {{overlapping accesses to 'x', but deinitialization requires exclusive access}}
         // expected-note @-2 {{conflicting access is here}}
         // expected-note @-3 {{used here}}
         // expected-note @-4 {{used here}}
@@ -975,7 +975,7 @@ func testLocalLetClosureCaptureConsuming(_ x: consuming SingleElt) {
         consumeVal(x) // expected-note {{consumed here}}
         // expected-note @-1 {{consumed again here}}
         borrowConsumeVal(x, x) // expected-note {{used here}}
-        // expected-error @-1 {{overlapping accesses, but deinitialization requires exclusive access}}
+        // expected-error @-1 {{overlapping accesses to 'x', but deinitialization requires exclusive access}}
         // expected-note @-2 {{conflicting access is here}}
         // expected-note @-3 {{consumed here}}
         // expected-note @-4 {{used here}}

test/SILOptimizer/allocbox_to_stack_ownership.sil

@@ -1435,3 +1435,19 @@ die:
   destroy_value %copy
   unreachable
 }
+
+// CHECK-LABEL: sil [ossa] @test_access_enforcement :
+// OPT:           begin_access [modify] [dynamic]
+// MANDATORY:     begin_access [modify] [static]
+// CHECK-LABEL: } // end sil function 'test_access_enforcement'
+sil [ossa] @test_access_enforcement : $(Int) -> Int {
+bb0(%0 : $Int):
+  %1 = alloc_box ${ var Int }
+  %2 = project_box %1 : ${ var Int }, 0
+  %3 = begin_access [modify] [dynamic] %2
+  store %0 to [trivial] %3
+  end_access %3
+  %4 = load [trivial] %2
+  destroy_value %1
+  return %4 : $Int
+}