Add the yk_unroll_safe attribute.

This function attribute tells the Yk JIT that it is OK to inline the
decorated function into the trace, even if it contains loops. By default
functions containing loops will not be inlined into traces by default
any more.

These new behaviours are enabled by passing
-fyk-noinline-funcs-with-loops to clang.

Author: vext01

clang/include/clang/Basic/Attr.td

@@ -4072,3 +4072,10 @@ def YkOutline : InheritableAttr {
   let Documentation = [YkOutlineDocs];
   let SimpleHandler = 1;
 }
+
+def YkUnrollSafe : InheritableAttr {
+  let Spellings = [GCC<"yk_unroll_safe">, Declspec<"yk_unroll_safe">];
+  let Subjects = SubjectList<[Function]>;
+  let Documentation = [YkUnrollSafeDocs];
+  let SimpleHandler = 1;
+}

clang/include/clang/Basic/AttrDocs.td

@@ -6667,3 +6667,12 @@ def YkOutlineDocs : Documentation {
   to the annotated function.
   }];
 }
+
+def YkUnrollSafeDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+  The ``yk_unroll_safe`` attribute tells the Yk JIT that during tracing, calls
+  to the annotated function can be inlined, even if it contains loops. By
+  default only calls to functions that contain no loops are inlined.
+  }];
+}

clang/include/clang/Basic/CodeGenOptions.def

@@ -333,6 +333,8 @@ CODEGENOPT(DebugColumnInfo, 1, 0) ///< Whether or not to use column information
 CODEGENOPT(DebugTypeExtRefs, 1, 0) ///< Whether or not debug info should contain
                                    ///< external references to a PCH or module.
 
+CODEGENOPT(YkNoinlineFuncsWithLoops, 1, 0) ///< Yk JIT: by default don't inline functions
+                                           ///< containing loops.
 CODEGENOPT(DebugExplicitImport, 1, 0)  ///< Whether or not debug info should
                                        ///< contain explicit imports for
                                        ///< anonymous namespaces

clang/include/clang/Driver/Options.td

@@ -4739,6 +4739,11 @@ def finline_limit_EQ : Joined<["-"], "finline-limit=">, Group<clang_ignored_gcc_
 defm finline_limit : BooleanFFlag<"inline-limit">, Group<clang_ignored_gcc_optimization_f_Group>;
 defm inline_small_functions : BooleanFFlag<"inline-small-functions">,
     Group<clang_ignored_gcc_optimization_f_Group>;
+defm yk_noinline_funcs_with_loops : BoolOption<"f", "yk-noinline-funcs-with-loops",
+  CodeGenOpts<"YkNoinlineFuncsWithLoops">, DefaultFalse,
+  NegFlag<SetFalse, []>,
+  PosFlag<SetTrue, [], "Don't inline functions containing loops by default (used for the Yk JIT).">,
+  BothFlags<[CC1Option, CoreOption]>>;
 defm ipa_cp : BooleanFFlag<"ipa-cp">,
     Group<clang_ignored_gcc_optimization_f_Group>;
 defm ivopts : BooleanFFlag<"ivopts">, Group<clang_ignored_gcc_optimization_f_Group>;

clang/lib/CodeGen/CodeGenModule.cpp

@@ -49,10 +49,12 @@
 #include "clang/Frontend/FrontendDiagnostic.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
@@ -2049,10 +2051,41 @@ void CodeGenModule::SetLLVMFunctionAttributesForDefinition(const Decl *D,
   }
 
   if (D->hasAttr<YkOutlineAttr>()) {
-    // Prevent the Yk trace compiler from inlining the call.
-    B.addAttribute("yk_outline");
-    // Prevent LLVM from inlining the call when optimising the trace.
     B.addAttribute(llvm::Attribute::NoInline);
+    B.addAttribute("yk_outline");
+  }
+
+  // Mark all functions containing loops with `yk_outline` unless the function
+  // was annotated `yk_unroll_safe`.
+  //
+  // This needs to be gated so that we don't impact the normal clang and LLVM
+  // tests.
+  //
+  // If `F->empty()` then this is merely a function declaration for which we
+  // have no IR. In this case the JIT will be unable to inline calls to this
+  // function anyway, so there's no need to conservatively add `yk_outline`.
+  if (CodeGenOpts.YkNoinlineFuncsWithLoops && !F->empty()) {
+    llvm::DominatorTree DT(*F);
+    llvm::LoopInfo LI(DT);
+    if (!LI.empty()) {
+      if (!D->hasAttr<YkUnrollSafeAttr>())
+        B.addAttribute("yk_outline");
+
+      // Note that we still mark the loopy function `F` with `noinline` (to
+      // block it being inlined during AOT compilation) regardless of if it can
+      // be inlined into the trace or not, because:
+      //
+      // * if `F` is to be inlined into a trace and we don't block AOT inlining
+      // then `F` may get AOT inlined into a parent function which contains
+      // loops but isn't marked `yk_unroll_safe`. The trace would have to call
+      // the parent, meaning that `F` would not be inlined into the trace!
+      //
+      // * if `F` is to be outlined in a trace and we don't block AOT inlining,
+      // then `F` may be inlined into a parent function and tracing the parent
+      // function would give a trace where `F` was actually inlined instead of
+      // outlined.
+      B.addAttribute(llvm::Attribute::NoInline);
+    }
   }
 
   F->addFnAttrs(B);

clang/lib/Driver/ToolChains/Clang.cpp

@@ -5399,6 +5399,10 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
   else if (UnwindTables)
     CmdArgs.push_back("-funwind-tables=1");
 
+  if (Args.hasFlag(options::OPT_fyk_noinline_funcs_with_loops,
+                   options::OPT_fno_yk_noinline_funcs_with_loops, false))
+    CmdArgs.push_back("-fyk-noinline-funcs-with-loops");
+
   // Prepare `-aux-target-cpu` and `-aux-target-feature` unless
   // `--gpu-use-aux-triple-only` is specified.
   if (!Args.getLastArg(options::OPT_gpu_use_aux_triple_only) &&

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -197,5 +197,6 @@
 // CHECK-NEXT: XRayInstrument (SubjectMatchRule_function, SubjectMatchRule_objc_method)
 // CHECK-NEXT: XRayLogArgs (SubjectMatchRule_function, SubjectMatchRule_objc_method)
 // CHECK-NEXT: YkOutline (SubjectMatchRule_function)
+// CHECK-NEXT: YkUnrollSafe (SubjectMatchRule_function)
 // CHECK-NEXT: ZeroCallUsedRegs (SubjectMatchRule_function)
 // CHECK-NEXT: End of supported attributes.

clang/test/Yk/yk_noinline_loopy_funcs.c

@@ -0,0 +1,15 @@
+// RUN: %clang -emit-llvm -fyk-noinline-funcs-with-loops -c -O0 -o - %s | llvm-dis | grep "call .*@never_aot_inline("
+// RUN: %clang -emit-llvm -fyk-noinline-funcs-with-loops -c -O1 -o - %s | llvm-dis | grep "call .*@never_aot_inline("
+// RUN: %clang -emit-llvm -fyk-noinline-funcs-with-loops -c -O2 -o - %s | llvm-dis | grep "call .*@never_aot_inline("
+// RUN: %clang -emit-llvm -fyk-noinline-funcs-with-loops -c -O3 -o - %s | llvm-dis | grep "call .*@never_aot_inline("
+
+#include <stdio.h>
+
+static void never_aot_inline(int i) {
+  while (i--)
+    putchar('.');
+}
+
+int main(int argc, char **argv) {
+  never_aot_inline(argc);
+}