IRGen: Add the ability to mark certain generic entry points in back t…

…races Mark generic function calls with concrete parameters, generic v-table calls and generic witness table calls where self is generic.
swiftlang · Oct 3, 2024 · 410bf3c · 410bf3c
1 parent 79a0afb
commit 410bf3c
Show file tree

Hide file tree

Showing 13 changed files with 287 additions and 10 deletions.
diff --git a/include/swift/AST/IRGenOptions.h b/include/swift/AST/IRGenOptions.h
@@ -484,6 +484,8 @@ class IRGenOptions {
   // (LLVM's 'frame-pointer=all').
   unsigned AsyncFramePointerAll : 1;
 
+  unsigned UseProfilingMarkerThunks : 1;
+
   /// The number of threads for multi-threaded code generation.
   unsigned NumThreads = 0;
 
@@ -573,6 +575,7 @@ class IRGenOptions {
         DisableReadonlyStaticObjects(false), CollocatedMetadataFunctions(false),
         ColocateTypeDescriptors(true), UseRelativeProtocolWitnessTables(false),
         UseFragileResilientProtocolWitnesses(false),
+        UseProfilingMarkerThunks(false),
         EnableHotColdSplit(false), EmitAsyncFramePushPopMetadata(false),
         AsyncFramePointerAll(false),
         CmdArgs(), SanitizeCoverage(llvm::SanitizerCoverageOptions()),

diff --git a/include/swift/Option/FrontendOptions.td b/include/swift/Option/FrontendOptions.td
@@ -1349,6 +1349,13 @@ def disable_new_llvm_pass_manager :
   Flag<["-"], "disable-new-llvm-pass-manager">,
   HelpText<"Disable the new llvm pass manager">;
 
+def enable_profiling_marker_thunks :
+  Flag<["-"], "enable-profiling-marker-thunks">,
+  HelpText<"Enable profiling marker thunks">;
+def disable_profiling_marker_thunks :
+  Flag<["-"], "disable-profiling-marker-thunks">,
+  HelpText<"Disable profiling marker thunks">;
+
 def enable_objective_c_protocol_symbolic_references :
   Flag<["-"], "enable-objective-c-protocol-symbolic-references">,
   HelpText<"Enable objective-c protocol symbolic references">;

diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -3469,6 +3469,10 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Args.hasFlag(OPT_enable_fragile_resilient_protocol_witnesses,
                  OPT_disable_fragile_resilient_protocol_witnesses,
                  Opts.UseFragileResilientProtocolWitnesses);
+  Opts.UseProfilingMarkerThunks =
+    Args.hasFlag(OPT_enable_profiling_marker_thunks,
+                 OPT_disable_profiling_marker_thunks,
+                 Opts.UseProfilingMarkerThunks);
   Opts.EnableHotColdSplit =
       Args.hasFlag(OPT_enable_split_cold_code,
                    OPT_disable_split_cold_code,

diff --git a/lib/IRGen/CallEmission.h b/lib/IRGen/CallEmission.h
@@ -77,6 +77,8 @@ class CallEmission {
   /// RemainingArgsForCallee, at least between calls.
   bool EmittedCall;
 
+  bool UseProfilingThunk = false;
+
   /// The basic block to which the call to a potentially throwing foreign
   /// function should jump to continue normal execution of the program.
   llvm::BasicBlock *invokeNormalDest = nullptr;
@@ -123,6 +125,10 @@ class CallEmission {
     return CurCallee.getSubstitutions();
   }
 
+  void useProfilingThunk() {
+    UseProfilingThunk = true;
+  }
+
   virtual void begin();
   virtual void end();
   virtual SILType getParameterType(unsigned index) = 0;

diff --git a/lib/IRGen/Callee.h b/lib/IRGen/Callee.h
@@ -374,6 +374,15 @@ namespace irgen {
     }
 
   public:
+
+    FunctionPointer withProfilingThunk(llvm::Function *thunk) const {
+      auto res = FunctionPointer(kind, thunk, nullptr/*secondaryValue*/,
+                                 AuthInfo, Sig);
+      res.useSignature = useSignature;
+      return res;
+    }
+
+
     FunctionPointer()
         : kind(FunctionPointer::Kind::Function), Value(nullptr),
           SecondaryValue(nullptr) {}

diff --git a/lib/IRGen/GenCall.cpp b/lib/IRGen/GenCall.cpp
@@ -52,6 +52,7 @@
 #include "GenType.h"
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
+#include "IRGenMangler.h"
 #include "LoadableTypeInfo.h"
 #include "NativeConventionSchema.h"
 #include "Signature.h"
@@ -3219,7 +3220,31 @@ llvm::CallBase *CallEmission::emitCallSite() {
     } else
       IGF.setCallsThunksWithForeignExceptionTraps();
   }
-  auto call = createCall(fn, Args);
+
+  auto fnToCall = fn;
+  if (UseProfilingThunk) {
+    assert(fnToCall.isConstant() && "Non constant function in profiling thunk");
+    auto genericFn = cast<llvm::Function>(fnToCall.getRawPointer());
+    auto replacementTypes = CurCallee.getSubstitutions().getReplacementTypes();
+    llvm::SmallString<64> name;
+    {
+      llvm::raw_svector_ostream os(name);
+      os << "__swift_prof_thunk__generic_func__";
+      os << replacementTypes.size();
+      os << "__";
+      for (auto replTy : replacementTypes) {
+        IRGenMangler mangler;
+        os << mangler.mangleTypeMetadataFull(replTy->getCanonicalType());
+        os << "___";
+      }
+      os << "fun__";
+    }
+    auto *thunk = IGF.IGM.getOrCreateProfilingThunk(genericFn, name);
+    fnToCall = fnToCall.withProfilingThunk(thunk);
+
+  }
+
+  auto call = createCall(fnToCall, Args);
   if (invokeNormalDest)
     IGF.Builder.emitBlock(invokeNormalDest);
 

diff --git a/lib/IRGen/GenDecl.cpp b/lib/IRGen/GenDecl.cpp
@@ -6278,3 +6278,62 @@ void IRGenModule::setColocateMetadataSection(llvm::Function *f) {
     break;
   }
 }
+
+llvm::Function *IRGenModule::getOrCreateProfilingThunk(
+  llvm::Function *f,
+  StringRef prefix) {
+
+  llvm::SmallString<32> name;
+  {
+    llvm::raw_svector_ostream os(name);
+    os << prefix;
+    os << f->getName();
+  }
+
+  auto thunk = cast<llvm::Function>(
+    getOrCreateHelperFunction(name, f->getReturnType(),
+                              f->getFunctionType()->params(),
+    [&](IRGenFunction &IGF) {
+      Explosion args = IGF.collectParameters();
+      auto res = IGF.Builder.CreateCall(f->getFunctionType(), f, args.getAll());
+      res->setAttributes(f->getAttributes());
+      (void)args.claimAll();
+      if (res->getType()->isVoidTy()) {
+        IGF.Builder.CreateRetVoid();
+      } else {
+        IGF.Builder.CreateRet(res);
+      }
+    }, /*isNoInline*/ true));
+
+  thunk->setAttributes(f->getAttributes());
+  thunk->setCallingConv(f->getCallingConv());
+  thunk->setDLLStorageClass(f->getDLLStorageClass());
+  if (f->getComdat())
+    thunk->setComdat(f->getParent()->getOrInsertComdat(thunk->getName()));
+  setMustHaveFramePointer(thunk);
+  thunk->addFnAttr(llvm::Attribute::NoInline);
+
+  return cast<llvm::Function>(thunk);
+}
+
+llvm::Function*
+IRGenModule::getAddrOfWitnessTableProfilingThunk(
+  llvm::Function *witness,
+  const NormalProtocolConformance &conformance) {
+
+  assert(
+    conformance.getDeclContext()->getSelfNominalTypeDecl()->isGenericContext());
+
+  return getOrCreateProfilingThunk(witness,
+                            "__swift_prof_thunk__generic_witness__");
+}
+
+llvm::Function *
+IRGenModule::getAddrOfVTableProfilingThunk(
+  llvm::Function *vTableFun, ClassDecl *classDecl) {
+
+  assert(classDecl->getSelfNominalTypeDecl()->isGenericContext());
+
+  return getOrCreateProfilingThunk(vTableFun,
+                            "__swift_prof_thunk__generic_vtable__");
+}
diff --git a/lib/IRGen/GenMeta.cpp b/lib/IRGen/GenMeta.cpp
@@ -287,7 +287,8 @@ static Flags getMethodDescriptorFlags(ValueDecl *fn) {
 static void buildMethodDescriptorFields(IRGenModule &IGM,
                              const SILVTable *VTable,
                              SILDeclRef fn,
-                             ConstantStructBuilder &descriptor) {
+                             ConstantStructBuilder &descriptor,
+                             ClassDecl *classDecl) {
   auto *func = cast<AbstractFunctionDecl>(fn.getDecl());
   // Classify the method.
   using Flags = MethodDescriptorFlags;
@@ -318,6 +319,13 @@ static void buildMethodDescriptorFields(IRGenModule &IGM,
       descriptor.addRelativeAddress(implFn);
     } else {
       llvm::Function *implFn = IGM.getAddrOfSILFunction(impl, NotForDefinition);
+
+      if (IGM.getOptions().UseProfilingMarkerThunks &&
+          classDecl->getSelfNominalTypeDecl()->isGenericContext() &&
+          !impl->getLoweredFunctionType()->isCoroutine()) {
+        implFn = IGM.getAddrOfVTableProfilingThunk(implFn, classDecl);
+      }
+
       descriptor.addCompactFunctionReference(implFn);
     }
   } else {
@@ -328,7 +336,8 @@ static void buildMethodDescriptorFields(IRGenModule &IGM,
 }
 
 void IRGenModule::emitNonoverriddenMethodDescriptor(const SILVTable *VTable,
-                                                    SILDeclRef declRef) {
+                                                    SILDeclRef declRef,
+                                                    ClassDecl *classDecl) {
   auto entity = LinkEntity::forMethodDescriptor(declRef);
   auto *var = cast<llvm::GlobalVariable>(
       getAddrOfLLVMVariable(entity, ConstantInit(), DebugTypeInfo()));
@@ -343,7 +352,7 @@ void IRGenModule::emitNonoverriddenMethodDescriptor(const SILVTable *VTable,
   ConstantInitBuilder ib(*this);
   ConstantStructBuilder sb(ib.beginStruct(MethodDescriptorStructTy));
 
-  buildMethodDescriptorFields(*this, VTable, declRef, sb);
+  buildMethodDescriptorFields(*this, VTable, declRef, sb, classDecl);
 
   auto init = sb.finishAndCreateFuture();
 
@@ -2081,7 +2090,7 @@ namespace {
 
       // Actually build the descriptor.
       auto descriptor = B.beginStruct(IGM.MethodDescriptorStructTy);
-      buildMethodDescriptorFields(IGM, VTable, fn, descriptor);
+      buildMethodDescriptorFields(IGM, VTable, fn, descriptor, getType());
       descriptor.finishAndAddTo(B);
 
       // Emit method dispatch thunk if the class is resilient.
@@ -2129,7 +2138,7 @@ namespace {
       // exist in the table in the class's context descriptor since it isn't
       // in the vtable, but external clients need to be able to link against the
       // symbol.
-      IGM.emitNonoverriddenMethodDescriptor(VTable, fn);
+      IGM.emitNonoverriddenMethodDescriptor(VTable, fn, getType());
     }
 
     void addOverrideTable() {

diff --git a/lib/IRGen/GenProto.cpp b/lib/IRGen/GenProto.cpp
@@ -1631,7 +1631,17 @@ class AccessorConformanceInfo : public ConformanceInfo {
         if (Func->isAsync()) {
           witness = IGM.getAddrOfAsyncFunctionPointer(Func);
         } else {
-          witness = IGM.getAddrOfSILFunction(Func, NotForDefinition);
+
+          auto *conformance = dyn_cast<NormalProtocolConformance>(&Conformance);
+          auto f = IGM.getAddrOfSILFunction(Func, NotForDefinition);
+          if (IGM.getOptions().UseProfilingMarkerThunks &&
+              conformance &&
+              conformance->getDeclContext()->
+                getSelfNominalTypeDecl()->isGenericContext() &&
+                !Func->getLoweredFunctionType()->isCoroutine())
+            witness = IGM.getAddrOfWitnessTableProfilingThunk(f, *conformance);
+          else witness = f;
+
         }
       } else {
         // The method is removed by dead method elimination.
@@ -1994,11 +2004,18 @@ void ResilientWitnessTableBuilder::collectResilientWitnesses(
 
     SILFunction *Func = entry.getMethodWitness().Witness;
     llvm::Constant *witness;
+    bool isGenericConformance =
+      conformance.getDeclContext()->getSelfNominalTypeDecl()->isGenericContext();
     if (Func) {
       if (Func->isAsync())
         witness = IGM.getAddrOfAsyncFunctionPointer(Func);
-      else
-        witness = IGM.getAddrOfSILFunction(Func, NotForDefinition);
+      else {
+        auto f = IGM.getAddrOfSILFunction(Func, NotForDefinition);
+        if (isGenericConformance && IGM.getOptions().UseProfilingMarkerThunks &&
+            !Func->getLoweredFunctionType()->isCoroutine())
+          witness = IGM.getAddrOfWitnessTableProfilingThunk(f, conformance);
+        else witness = f;
+      }
     } else {
       // The method is removed by dead method elimination.
       // It should be never called. We add a null pointer.

diff --git a/lib/IRGen/IRGenModule.cpp b/lib/IRGen/IRGenModule.cpp
@@ -1477,6 +1477,12 @@ void IRGenModule::setHasNoFramePointer(llvm::Function *F) {
   F->addFnAttrs(b);
 }
 
+void IRGenModule::setMustHaveFramePointer(llvm::Function *F) {
+  llvm::AttrBuilder b(getLLVMContext());
+  b.addAttribute("frame-pointer", "all");
+  F->addFnAttrs(b);
+}
+
 /// Construct initial function attributes from options.
 void IRGenModule::constructInitialFnAttributes(
     llvm::AttrBuilder &Attrs, OptimizationMode FuncOptMode,

diff --git a/lib/IRGen/IRGenModule.h b/lib/IRGen/IRGenModule.h
@@ -1586,6 +1586,7 @@ private:                                                                       \
       StackProtectorMode stackProtect = StackProtectorMode::NoStackProtector);
   void setHasNoFramePointer(llvm::AttrBuilder &Attrs);
   void setHasNoFramePointer(llvm::Function *F);
+  void setMustHaveFramePointer(llvm::Function *F);
   llvm::AttributeList constructInitialAttributes();
   StackProtectorMode shouldEmitStackProtector(SILFunction *f);
 
@@ -1660,7 +1661,8 @@ private:                                                                       \
   llvm::Constant *getAddrOfMethodDescriptor(SILDeclRef declRef,
                                             ForDefinition_t forDefinition);
   void emitNonoverriddenMethodDescriptor(const SILVTable *VTable,
-                                         SILDeclRef declRef);
+                                         SILDeclRef declRef,
+                                         ClassDecl *classDecl);
 
   Address getAddrOfEnumCase(EnumElementDecl *Case,
                             ForDefinition_t forDefinition);
@@ -1821,6 +1823,16 @@ private:                                                                       \
                        bool isDynamicallyReplaceableImplementation = false,
                        bool shouldCallPreviousImplementation = false);
 
+  llvm::Function *
+  getAddrOfWitnessTableProfilingThunk(llvm::Function *witness,
+                                      const NormalProtocolConformance &C);
+
+  llvm::Function *
+  getAddrOfVTableProfilingThunk(llvm::Function *f, ClassDecl *decl);
+
+  llvm::Function *getOrCreateProfilingThunk(llvm::Function *f,
+                                            StringRef prefix);
+
   void emitDynamicReplacementOriginalFunctionThunk(SILFunction *f);
 
   llvm::Function *getAddrOfContinuationPrototype(CanSILFunctionType fnType);

diff --git a/lib/IRGen/IRGenSIL.cpp b/lib/IRGen/IRGenSIL.cpp
@@ -3819,6 +3819,15 @@ void IRGenSILFunction::visitFullApplySite(FullApplySite site) {
     SubstitutionMap subMap = site.getSubstitutionMap();
     emitPolymorphicArguments(*this, origCalleeType,
                              subMap, &witnessMetadata, llArgs);
+
+    // We currently only support non-async calls with profiling thunks.
+    if (IGM.getOptions().UseProfilingMarkerThunks &&
+        isa<FunctionRefInst>(site.getCallee()) &&
+        !site.getOrigCalleeType()->isAsync() &&
+        subMap.hasAnySubstitutableParams() &&
+        !subMap.hasArchetypes()) {
+      emission->useProfilingThunk();
+    }
   }
 
   if (calleeFP.shouldPassContinuationDirectly()) {