[assembly-vision] If a nominal type is marked with @_assemblyVision emit AssemblyVisionRemarks for all of the nominal type's methods.

include/swift/AST/Attr.def

@@ -673,7 +673,7 @@ SIMPLE_DECL_ATTR(_distributedActorIndependent, DistributedActorIndependent,
   119)
 
 SIMPLE_DECL_ATTR(_assemblyVision, EmitAssemblyVisionRemarks,
-  OnFunc | UserInaccessible | NotSerialized |
+  OnFunc | UserInaccessible | NotSerialized | OnNominalType |
   ABIStableToAdd | ABIStableToRemove | APIStableToAdd | APIStableToRemove,
   120)
 

include/swift/AST/Decl.h

@@ -3279,6 +3279,12 @@ class NominalTypeDecl : public GenericTypeDecl, public IterableDeclContext {
     return getAttrs().hasSemanticsAttr(attrValue);
   }
 
+  /// Returns true if we should emit assembly vision remarks on all methods of
+  /// this nominal type.
+  bool shouldEmitAssemblyVisionRemarksOnMethods() const {
+    return getAttrs().hasAttribute<EmitAssemblyVisionRemarksAttr>();
+  }
+
   /// Whether this declaration has a synthesized memberwise initializer.
   bool hasMemberwiseInitializer() const;
 

lib/SIL/Utils/OptimizationRemark.cpp

@@ -135,15 +135,28 @@ static bool hasForceEmitSemanticAttr(SILFunction &fn, StringRef passName) {
   });
 }
 
+static bool isMethodWithForceEmitSemanticAttrNominalType(SILFunction &fn) {
+  if (!fn.hasSelfParam())
+    return false;
+
+  auto selfType = fn.getSelfArgument()->getType();
+  auto *nomType = selfType.getNominalOrBoundGenericNominal();
+  if (!nomType)
+    return false;
+  return nomType->shouldEmitAssemblyVisionRemarksOnMethods();
+}
+
 Emitter::Emitter(StringRef passName, SILFunction &fn)
     : fn(fn), passName(passName),
       passedEnabled(
           hasForceEmitSemanticAttr(fn, passName) ||
+          isMethodWithForceEmitSemanticAttrNominalType(fn) ||
           (fn.getASTContext().LangOpts.OptimizationRemarkPassedPattern &&
            fn.getASTContext().LangOpts.OptimizationRemarkPassedPattern->match(
                passName))),
       missedEnabled(
           hasForceEmitSemanticAttr(fn, passName) ||
+          isMethodWithForceEmitSemanticAttrNominalType(fn) ||
           (fn.getASTContext().LangOpts.OptimizationRemarkMissedPattern &&
            fn.getASTContext().LangOpts.OptimizationRemarkMissedPattern->match(
                passName))) {}

lib/SILOptimizer/Transforms/AssemblyVisionRemarkGenerator.cpp

@@ -773,11 +773,35 @@ class AssemblyVisionRemarkGenerator : public SILFunctionTransform {
     // TODO: Put this on LangOpts as a helper.
     auto &langOpts = fn->getASTContext().LangOpts;
 
-    return bool(langOpts.OptimizationRemarkMissedPattern) ||
-           bool(langOpts.OptimizationRemarkPassedPattern) ||
-           fn->getModule().getSILRemarkStreamer() ||
-           fn->hasSemanticsAttrThatStartsWith(
-               semantics::FORCE_EMIT_OPT_REMARK_PREFIX);
+    // If we are supposed to emit remarks, always emit.
+    if (bool(langOpts.OptimizationRemarkMissedPattern) ||
+        bool(langOpts.OptimizationRemarkPassedPattern) ||
+        fn->getModule().getSILRemarkStreamer())
+      return true;
+
+    // Otherwise, first check if our function has a force emit opt remark prefix
+    // semantics tag.
+    if (fn->hasSemanticsAttrThatStartsWith(
+            semantics::FORCE_EMIT_OPT_REMARK_PREFIX))
+      return true;
+
+    // Otherwise, check if we have a self parameter that is a nominal type that
+    // is marked with the @_assemblyVision attribute.
+    if (fn->hasSelfParam()) {
+      if (auto *nomType = fn->getSelfArgument()
+                              ->getType()
+                              .getNominalOrBoundGenericNominal()) {
+        LLVM_DEBUG(llvm::dbgs() << "Checking for remark on: "
+                                << nomType->getName().get() << "\n");
+        if (nomType->shouldEmitAssemblyVisionRemarksOnMethods()) {
+          LLVM_DEBUG(llvm::dbgs() << "Success! Will emit remarks!!\n");
+          return true;
+        }
+        LLVM_DEBUG(llvm::dbgs() << "Fail! No remarks will be emitted!!\n");
+      }
+    }
+
+    return false;
   }
 
   /// The entry point to the transformation.
@@ -791,14 +815,24 @@ class AssemblyVisionRemarkGenerator : public SILFunctionTransform {
     // unless we were asked by the user to emit them.
     if (!ForceVisitImplicitAutogeneratedFunctions) {
       // Skip implicit functions generated by Sema.
-      if (auto *ctx = fn->getDeclContext())
-        if (auto *decl = ctx->getAsDecl())
-          if (decl->isImplicit())
+      if (auto *ctx = fn->getDeclContext()) {
+        if (auto *decl = ctx->getAsDecl()) {
+          if (decl->isImplicit()) {
+            LLVM_DEBUG(llvm::dbgs() << "Skipping implicit decl function: "
+                                    << fn->getName() << "\n");
             return;
+          }
+        }
+      }
+
       // Skip autogenerated functions generated by SILGen.
-      if (auto loc = fn->getDebugScope()->getLoc())
-        if (loc.isAutoGenerated())
+      if (auto loc = fn->getDebugScope()->getLoc()) {
+        if (loc.isAutoGenerated()) {
+          LLVM_DEBUG(llvm::dbgs() << "Skipping autogenerated function: "
+                                  << fn->getName() << "\n");
           return;
+        }
+      }
     }
 
     LLVM_DEBUG(llvm::dbgs() << "Visiting: " << fn->getName() << "\n");

test/SILOptimizer/assemblyvision_remark/nominal_type_attributes.swift

@@ -0,0 +1,51 @@
+// RUN: %target-swiftc_driver -Osize -emit-sil %s -o /dev/null -Xfrontend -verify
+// REQUIRES: optimized_stdlib,swift_stdlib_no_asserts
+
+// Make sure we emit remarks on nominal types
+
+@inline(never)
+func callPrint(_ s: String) { print(s) }
+
+var global: String = "123"
+
+@_assemblyVision
+struct Struct {
+    func printMe() {
+        callPrint(global) // expected-remark {{begin exclusive access to value of type '}}
+                          // expected-note @-6 {{of 'global'}}
+                          // expected-remark @-2 {{end exclusive access to value of type '}}
+                          // expected-note @-8 {{of 'global'}}
+                          // expected-remark @-4 {{retain of type '}}
+                          // expected-note @-10 {{of 'global'}}
+                          // expected-remark @-6 {{release of type '}}
+                          // expected-note @-12 {{of 'global}}
+    }
+}
+
+// Negative test
+struct Struct2 {
+    func callPrintMe() {
+        callPrint(global)
+    }
+}
+
+@_assemblyVision
+enum Enum {
+    func callPrintMe() {
+        callPrint(global) // expected-remark {{begin exclusive access to value of type '}}
+                          // expected-note @-27 {{of 'global'}}
+                          // expected-remark @-2 {{end exclusive access to value of type '}}
+                          // expected-note @-29 {{of 'global'}}
+                          // expected-remark @-4 {{retain of type '}}
+                          // expected-note @-31 {{of 'global'}}
+                          // expected-remark @-6 {{release of type '}}
+                          // expected-note @-33 {{of 'global}}
+    }
+}
+
+// Negative test
+enum Enum2 {
+    func callPrintMe() {
+        callPrint("I am callPrinting 1")
+    }
+}