Merge pull request #20522 from shajrawi/block_args

[Exclusivity] Support exclusivity checks originating from block arguments

Author: Joe Shajrawi

lib/SIL/MemAccessUtils.cpp

@@ -287,15 +287,31 @@ AccessedStorage swift::findAccessedStorage(SILValue sourceAddr) {
         return AccessedStorage(address, AccessedStorage::Unidentified);
       return AccessedStorage();
 
-    // A block argument may be a box value projected out of
-    // switch_enum. Address-type block arguments are not allowed.
-    case ValueKind::SILPhiArgument:
+    case ValueKind::SILPhiArgument: {
+      auto *phiArg = cast<SILPhiArgument>(address);
+      bool allValsMatch = true;
+      SmallVector<SILValue, 8> incomingPhis;
+      phiArg->getIncomingPhiValues(incomingPhis);
+      if (!incomingPhis.empty()) {
+        auto firstVal = incomingPhis.front();
+        for (auto val : incomingPhis) {
+          if (val != firstVal) {
+            allValsMatch = false;
+            break;
+          }
+        }
+        if (allValsMatch) {
+          return findAccessedStorage(firstVal);
+        }
+      }
+      // A block argument may be a box value projected out of
+      // switch_enum. Address-type block arguments are not allowed.
       if (address->getType().isAddress())
         return AccessedStorage();
 
       checkSwitchEnumBlockArg(cast<SILPhiArgument>(address));
       return AccessedStorage(address, AccessedStorage::Unidentified);
-
+    }
     // Load a box from an indirect payload of an opaque enum.
     // We must have peeked past the project_box earlier in this loop.
     // (the indirectness makes it a box, the load is for address-only).

lib/SILOptimizer/Transforms/AccessEnforcementOpts.cpp

@@ -241,6 +241,13 @@ class AccessConflictAndMergeAnalysis {
     /// Map each begin access to its AccessInfo with index, data, and flags.
     /// Iterating over this map is nondeterministic. If it is necessary to order
     /// the accesses, then AccessInfo::getAccessIndex() can be used.
+    /// This maps contains every dynamic begin_access instruction,
+    /// even those with invalid storage:
+    /// We would like to keep track of unrecognized or invalid storage locations
+    /// Because they affect our decisions for recognized locations,
+    /// be it nested conflict or merging out of scope accesses.
+    /// The access map is just a “cache” of accesses.
+    /// Keeping those invalid ones just makes the lookup faster
     AccessMap accessMap;
 
     /// Instruction pairs we can merge the scope of
@@ -536,8 +543,6 @@ void AccessConflictAndMergeAnalysis::identifyBeginAccesses() {
       // here as an invalid `storage` value.
       const AccessedStorage &storage =
           findAccessedStorageNonNested(beginAccess->getSource());
-      if (!storage)
-        continue;
 
       auto iterAndSuccess = result.accessMap.try_emplace(
           beginAccess, static_cast<const AccessInfo &>(storage));

test/SILOptimizer/access_enforcement_opts.sil

@@ -1515,4 +1515,43 @@ bb0:
   return %10 : $()
 }
 
+// public func testPhiArgs() {
+// Check that we can merge scopes with Phi args - avoiding a crash we've seen in the past
+//
+// CHECK-LABEL: sil @testPhiArgs : $@convention(thin) () -> () {
+// CHECK: [[GLOBAL:%.*]] = global_addr @globalX : $*X
+// CHECK: cond_br {{.*}}, bb1, bb2
+// CHECK: bb1
+// CHECK: br bb3([[GLOBAL]] : $*X)
+// CHECK: bb2
+// CHECK: br bb3([[GLOBAL]] : $*X)
+// CHECK: bb3([[GLOBALPHI:%.*]] : $*X):
+// CHECK: [[BEGIN:%.*]] = begin_access [read] [dynamic] [[GLOBALPHI]] : $*X
+// CHECK-NEXT: load
+// CHECK-NEXT: load
+// CHECK-NEXT: end_access [[BEGIN]] : $*X
+// CHECK-NOT: begin_access
+// CHECK-LABEL: } // end sil function 'testPhiArgs'
+sil @testPhiArgs : $@convention(thin) () -> () {
+bb0:
+  %0 = global_addr @globalX: $*X
+  %cond = integer_literal $Builtin.Int1, 1
+  cond_br %cond, bb1, bb2
+  
+bb1:
+  br bb3(%0 : $*X)
+  
+bb2:
+  br bb3(%0 : $*X)
+  
+bb3(%1 : $*X):
+  %2 = begin_access [read] [dynamic] %1 : $*X
+  %3 = load %2 : $*X
+  end_access %2 : $*X
+  %4 = begin_access [read] [dynamic] %0 : $*X
+  %5 = load %4 : $*X
+  end_access %4 : $*X
+  %10 = tuple ()
+  return %10 : $()
+}