[assembly-vision] Change release to do backwards then forwards when inferring source locs.

TLDR: I fixed a whole in the assembly-vision opt-remark pass where we were not
emitting a remark for end of scope instructions at the beginning of blocks. Now
all of these instructions (strong_release, end_access) should always reliably
have a remark emitted for them.

----

I think that this is a pragmatic first solution to the problem of
strong_release, release_value being the first instruction of a block. For those
who are unaware, this issue is that for a long time we have searched backwards
first for "end of scope" like instructions. This then allows us to identify the
"end of scope" instruction as happening at the end of the previous statement
which is where the developer thinks it should be:

```
var global: Klass
func bar() -> @owned Klass { global }
func foo() {
   // We want the remark for the
   bar()                          // expected-remark {{retain}}
}
```

This makes sense since we want to show end of scope instructions as being
applied to the earlier code whose scope it is ending. We can be clear that it is
at the end of the statement by placing the carrot on the end of statement
SourceLoc so there isn't any confusion upon whether or not

That generally has delivered nice looking results, but what if our release is
the first instruction in the block? In that case, we do not have any instruction
that we can immediately use, so traditionally we just gave up and didn't emit
anything. This is not an acceptable solution! We should be able to emit
something for every retain/release in the program if we want users to be able to
rely upon this! Thus we need to be able to get source location information from
somewhere around

First before we begin, my approach here is informed by my seeing over time that
the optimizer does a pretty good job of not breaking SourceLoc info for
terminators.

With that in mind, there are two possible approaches here: using the terminator
from the previous block and searching forward at worst taking the SourceLoc of
the current block's terminator (or earlier if we find a good SourceLoc). I
wasn't sure what the correct thing to do was at the time so I didn't fix the
issue. After some thought, I realized that the correct solution is to if we fail
a backwards search, search forwards. The reason why is that since our remarks
runs late in the optimization pipeline, there is a very high likelihood that if
we aren't folded into our previous block that there is a true need in the
program for conditional control flow here. We want to avoid placing the release
out of such pieces of code since it is misleading to the user:

```

In this example there is a release inside the case for .x but none for .y. In
that case it is possible that we get a release for .f since payload is passed in
at +1 at SILGen time. In such a case, using the terminator of the previous block
would mean that we would have the release be marked as on payload instead of
inside the case of .x. By using the terminator of the releases block, we can

switch payload {
case let .x(f):
  ...
case let .y:
  ...
}
```

So using the terminator from the previous block would be
  misleading to the user. Instead it is better to pick a location after the release that way
  we know at least the instruction we are inferring from must in some sense be

With this fix, we should not emit locations for all retains, releases. We may
not identify a good source loc for all of them, but we will identify them.

optimization pipeline, if our block was not folded into the previous block there
is a very high liklihood that there is some sort of conditional control flow
that is truly necessary in the program. If we

this
generally implies that there is a real side effect in the program that is
requiring conditional code execution (since the optimizer would have folded it).

The reason why is that we are at least going to hit a terminator or a
side-effect having instruction that generally have debug info preserved by the
optimizer.

Author: gottesmm

benchmark/single-source/ChaCha.swift

@@ -361,6 +361,7 @@ func checkResult(_ plaintext: [UInt8]) {
 }
 
 @inline(never)
+@_assemblyVision
 public func run_ChaCha(_ N: Int) {
   let key = Array(repeating: UInt8(1), count: 32)
   let nonce = Array(repeating: UInt8(2), count: 12)

include/swift/SIL/OptimizationRemark.h

@@ -153,6 +153,7 @@ enum class SourceLocInferenceBehavior : unsigned {
   BackwardsThenForwards = BackwardScan | ForwardScan2nd,
   ForwardScanAlwaysInfer = ForwardScan | AlwaysInfer,
   BackwardScanAlwaysInfer = BackwardScan | AlwaysInfer,
+  BackwardThenForwardAlwaysInfer = BackwardScan | ForwardScan2nd | AlwaysInfer,
 };
 
 inline SourceLocInferenceBehavior operator&(SourceLocInferenceBehavior lhs,

lib/SIL/Utils/OptimizationRemark.cpp

@@ -161,6 +161,19 @@ static SourceLoc getLocForPresentation(SILLocation loc,
   llvm_unreachable("covered switch");
 }
 
+static bool instHasInferrableLoc(SILInstruction &inst) {
+  if (isa<DeallocStackInst>(inst) || isa<StrongRetainInst>(inst) ||
+      isa<StrongReleaseInst>(inst) || isa<RetainValueInst>(inst) ||
+      isa<ReleaseValueInst>(inst) || isa<EndAccessInst>(inst))
+    return false;
+
+  // Ignore empty tuples
+  if (auto *tup = dyn_cast<TupleInst>(&inst))
+    return tup->getNumOperands() != 0;
+
+  return true;
+}
+
 /// The user has passed us an instruction that for some reason has a source loc
 /// that can not be used. Search down the current block for an instruction with
 /// a valid source loc and use that instead.
@@ -169,11 +182,16 @@ inferOptRemarkSearchForwards(SILInstruction &i,
                              SourceLocPresentationKind presentationKind) {
   for (auto &inst :
        llvm::make_range(std::next(i.getIterator()), i.getParent()->end())) {
+    if (!instHasInferrableLoc(inst))
+      continue;
+    // Skip instructions without a loc we care about since we move it around.
     auto newLoc = getLocForPresentation(inst.getLoc(), presentationKind);
     if (auto inlinedLoc = inst.getDebugScope()->getOutermostInlineLocation())
       newLoc = getLocForPresentation(inlinedLoc, presentationKind);
-    if (newLoc.isValid())
+    if (newLoc.isValid()) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferring loc from " << inst);
       return newLoc;
+    }
   }
 
   return SourceLoc();
@@ -188,11 +206,15 @@ inferOptRemarkSearchBackwards(SILInstruction &i,
                               SourceLocPresentationKind presentationKind) {
   for (auto &inst : llvm::make_range(std::next(i.getReverseIterator()),
                                      i.getParent()->rend())) {
+    if (!instHasInferrableLoc(inst))
+      continue;
     auto loc = inst.getLoc();
     if (auto inlinedLoc = inst.getDebugScope()->getOutermostInlineLocation())
       loc = inlinedLoc;
-    if (auto result = getLocForPresentation(loc, presentationKind))
+    if (auto result = getLocForPresentation(loc, presentationKind)) {
+      LLVM_DEBUG(llvm::dbgs() << "Inferring loc from " << inst);
       return result;
+    }
   }
 
   return SourceLoc();

lib/SILOptimizer/Transforms/AssemblyVisionRemarkGenerator.cpp

@@ -604,7 +604,8 @@ void AssemblyVisionRemarkGeneratorInstructionVisitor::visitEndAccessInst(
     // Use the actual source loc of the begin_access if it works. Otherwise,
     // scan backwards.
     auto remark =
-        RemarkMissed("memory", *eai, SourceLocInferenceBehavior::BackwardScan,
+        RemarkMissed("memory", *eai,
+                     SourceLocInferenceBehavior::BackwardThenForwardAlwaysInfer,
                      SourceLocPresentationKind::EndRange)
         << "end exclusive access to value of type '"
         << NV("ValueType", eai->getOperand()->getType()) << "'";
@@ -649,7 +650,7 @@ void AssemblyVisionRemarkGeneratorInstructionVisitor::visitStrongReleaseInst(
 
     auto remark =
         RemarkMissed("memory", *sri,
-                     SourceLocInferenceBehavior::BackwardScanAlwaysInfer,
+                     SourceLocInferenceBehavior::BackwardThenForwardAlwaysInfer,
                      SourceLocPresentationKind::EndRange)
         << "release of type '" << NV("ValueType", sri->getOperand()->getType())
         << "'";
@@ -693,7 +694,7 @@ void AssemblyVisionRemarkGeneratorInstructionVisitor::visitReleaseValueInst(
     // Releases end a lifetime scope so we infer scan backward.
     auto remark =
         RemarkMissed("memory", *rvi,
-                     SourceLocInferenceBehavior::BackwardScanAlwaysInfer)
+                     SourceLocInferenceBehavior::BackwardThenForwardAlwaysInfer)
         << "release of type '" << NV("ValueType", rvi->getOperand()->getType())
         << "'";
     for (auto arg : inferredArgs) {

test/SILOptimizer/assemblyvision_remark/basic.swift

@@ -15,6 +15,8 @@ public func getGlobal() -> Klass {
                   // expected-note @-5:12 {{of 'global'}}
                   // expected-remark @-2:12 {{begin exclusive access to value of type 'Klass'}}
                   // expected-note @-7:12 {{of 'global'}}
+                  // expected-remark @-4:12 {{end exclusive access to value of type 'Klass'}}
+                  // expected-note @-9:12 {{of 'global'}}
 }
 
 // Make sure that the retain msg is at the beginning of the print and the

test/SILOptimizer/assemblyvision_remark/basic_yaml.swift

@@ -28,7 +28,7 @@ public var global = Klass() // expected-remark {{heap allocated ref of type 'Kla
 // CHECK-NEXT: Pass:            sil-assembly-vision-remark-gen
 // CHECK-NEXT: Name:            sil.memory
 // CHECK-NEXT: DebugLoc:        { File: '{{.*}}basic_yaml.swift', 
-// CHECK-NEXT:                    Line: [[# @LINE + 27 ]], Column: 12 }
+// CHECK-NEXT:                    Line: [[# @LINE + 42 ]], Column: 12 }
 // CHECK-NEXT: Function:        'getGlobal()'
 // CHECK-NEXT: Args:
 // CHECK-NEXT:   - String:          'begin exclusive access to value of type '''
@@ -42,6 +42,21 @@ public var global = Klass() // expected-remark {{heap allocated ref of type 'Kla
 // CHECK: --- !Missed
 // CHECK-NEXT: Pass:            sil-assembly-vision-remark-gen
 // CHECK-NEXT: Name:            sil.memory
+// CHECK-NEXT: DebugLoc:        { File: '{{.*}}basic_yaml.swift', 
+// CHECK-NEXT:                    Line: [[# @LINE + 27 ]], Column: 12 }
+// CHECK-NEXT: Function:        'getGlobal()'
+// CHECK-NEXT: Args:
+// CHECK-NEXT:   - String:          'end exclusive access to value of type '''
+// CHECK-NEXT:   - ValueType:       Klass
+// CHECK-NEXT:   - String:          ''''
+// CHECK-NEXT:   - InferredValue:   'of ''global'''
+// CHECK-NEXT:     DebugLoc:        { File: '{{.*}}basic_yaml.swift', 
+// CHECK-NEXT:                        Line: [[# @LINE - 29 ]], Column: 12 }
+// CHECK-NEXT: ...
+//
+// CHECK: --- !Missed
+// CHECK-NEXT: Pass:            sil-assembly-vision-remark-gen
+// CHECK-NEXT: Name:            sil.memory
 // CHECK-NEXT: DebugLoc:        { File: '{{.*}}basic_yaml.swift',
 // CHECK-NEXT:                    Line: [[# @LINE + 12]], Column: 5 }
 // CHECK-NEXT: Function:        'getGlobal()'
@@ -51,14 +66,17 @@ public var global = Klass() // expected-remark {{heap allocated ref of type 'Kla
 // CHECK-NEXT:   - String:          ''''
 // CHECK-NEXT:   - InferredValue:   'of ''global'''
 // CHECK-NEXT:     DebugLoc:        { File: '{{.*}}basic_yaml.swift',
-// CHECK-NEXT:                        Line: [[# @LINE - 29 ]], Column: 12 }
+// CHECK-NEXT:                        Line: [[# @LINE - 44 ]], Column: 12 }
 // CHECK-NEXT: ...
 @inline(never)
 public func getGlobal() -> Klass {
     return global // expected-remark @:5 {{retain of type 'Klass'}}
-                  // expected-note @-34:12 {{of 'global'}}
+                  // expected-note @-49:12 {{of 'global'}}
                   // expected-remark @-2 {{begin exclusive access to value of type 'Klass'}}
-                  // expected-note @-36:12 {{of 'global'}}
+                  // expected-note @-51:12 {{of 'global'}}
+                  // NOTE: We really want the end access at :18, not :12. TODO Fix this!
+                  // expected-remark @-5 {{end exclusive access to value of type 'Klass'}}
+                  // expected-note @-54:12 {{of 'global'}}
 }
 
 // CHECK: --- !Missed

test/SILOptimizer/assemblyvision_remark/chacha.swift

@@ -0,0 +1,49 @@
+// RUN: %target-swiftc_driver -Osize -emit-sil %s -o /dev/null -Xfrontend -verify
+// REQUIRES: optimized_stdlib,swift_stdlib_no_asserts
+
+// An extraction from the benchmark ChaCha20 that we were not ignoring
+// dealloc_stack and other end scope instructions.
+
+enum ChaCha20 { }
+
+extension ChaCha20 {
+    @inline(never)
+    public static func encrypt<Key: Collection, Nonce: Collection, Bytes: MutableCollection>(bytes: inout Bytes, key: Key, nonce: Nonce, initialCounter: UInt32 = 0) where Bytes.Element == UInt8, Key.Element == UInt8, Nonce.Element == UInt8 {
+        print("I am lost...")
+    }
+}
+
+@inline(never)
+func checkResult(_ plaintext: [UInt8]) {
+    precondition(plaintext.first! == 6 && plaintext.last! == 254)
+    var hash: UInt64 = 0
+    for byte in plaintext {
+        // rotate
+        hash = (hash &<< 8) | (hash &>> (64 - 8))
+        hash ^= UInt64(byte)
+    }
+    precondition(hash == 0xa1bcdb217d8d14e4)
+}
+
+@_semantics("optremark.sil-assembly-vision-remark-gen")
+public func run_ChaCha(_ N: Int) {
+  let key = Array(repeating: UInt8(1), count: 32)
+  let nonce = Array(repeating: UInt8(2), count: 12)
+
+  var checkedtext = Array(repeating: UInt8(0), count: 1024)
+  ChaCha20.encrypt(bytes: &checkedtext, key: key, nonce: nonce)
+  checkResult(checkedtext) // expected-remark {{release of type '}}
+                           // expected-note @-3 {{of 'checkedtext}}
+
+  var plaintext = Array(repeating: UInt8(0), count: 30720)
+  for _ in 1...N {
+    ChaCha20.encrypt(bytes: &plaintext, key: key, nonce: nonce)
+    print(plaintext.first!) // expected-remark @:11 {{heap allocated ref of type '}}
+                            // expected-remark @-1:27 {{release of type '}}
+  }
+} // expected-remark {{release of type '}}
+  // expected-note @-7 {{of 'plaintext}}
+  // expected-remark @-2 {{release of type '}}
+  // expected-note @-16 {{of 'nonce}}
+  // expected-remark @-4 {{release of type '}}
+  // expected-note @-19 {{of 'key}}