[OpenMP] Introduce '-f[no-]openmp-target-jit' flag to control JIT for offloading

[jhuber6]

Summary

JIT support for OpenMP offloading was introduced in D139287. This patch
adds a simple flag that enables this mode. It simply requires enabling
-foffload-lto mode and --embed-bitcode in the linker wrapper. This
option implies LTO if it is not enabled.

[jdoerfert]

clang/lib/Driver/Driver.cpp

730

Should we overwrite unconditionally or warn/error if the combination is nonsensical?

[jhuber6]

clang/lib/Driver/Driver.cpp

730

Yeah, we might need something like that. Like if the user specifies thin then the monolithic linking won't work. Probably will need a check for that.

[jhuber6]

Adding an error print and some tests for it. We error if the user specified any -foffload-lto optoins that don't result in full LT.O

[jdoerfert]

LG

[gribozavr2]

The newly added test fails when tools/clang/include/clang/Config/config.h is configured with #define CLANG_DEFAULT_OPENMP_RUNTIME "libgomp". Please fix. Thanks!

[jhuber6]

In D141158#4035875, @gribozavr2 wrote:

The newly added test fails when tools/clang/include/clang/Config/config.h is configured with #define CLANG_DEFAULT_OPENMP_RUNTIME "libgomp". Please fix. Thanks!

Should be fixed now, let me know if there's anything else failing.

[gribozavr2]

Thank you!

clang/include/clang/Basic/DiagnosticDriverKinds.td

 def err_drv_invalid_rtlib_name : Error<
   "invalid runtime library name in argument '%0'">;
 def err_drv_unsupported_rtlib_for_platform : Error<
   "unsupported runtime library '%0' for platform '%1'">;
 def err_drv_invalid_unwindlib_name : Error<
   "invalid unwind library name in argument '%0'">;
 def err_drv_incompatible_unwindlib : Error<
   "--rtlib=libgcc requires --unwindlib=libgcc">;
+def err_drv_incompatible_options : Error<
+  "The combination of '%0' and '%1' is incompatible">;
 def err_drv_invalid_stdlib_name : Error<
   "invalid library name in argument '%0'">;
 def err_drv_invalid_output_with_multiple_archs : Error<
   "cannot use '%0' output with multiple -arch options">;
 def err_drv_no_input_files : Error<"no input files">;
 def err_drv_output_argument_with_multiple_files : Error<
   "cannot specify -o when generating multiple output files">;
 def err_drv_out_file_argument_with_multiple_sources : Error<

clang/include/clang/Driver/Options.td

 def fopenmp_assume_no_nested_parallelism : Flag<["-"], "fopenmp-assume-no-nested-parallelism">, Group<f_Group>,
   Flags<[CC1Option, NoArgumentUnused, HelpHidden]>,
   HelpText<"Assert no nested parallel regions in the GPU">,
   MarshallingInfoFlag<LangOpts<"OpenMPNoNestedParallelism">>;
 def fopenmp_offload_mandatory : Flag<["-"], "fopenmp-offload-mandatory">, Group<f_Group>,
   Flags<[CC1Option, NoArgumentUnused]>,
   HelpText<"Do not create a host fallback if offloading to the device fails.">,
   MarshallingInfoFlag<LangOpts<"OpenMPOffloadMandatory">>;
+def fopenmp_target_jit : Flag<["-"], "fopenmp-target-jit">, Group<f_Group>,
+  Flags<[CoreOption, NoArgumentUnused]>,
+  HelpText<"Emit code that can be JIT compiled for OpenMP offloading. Implies -foffload-lto=full">;
+def fno_openmp_target_jit : Flag<["-"], "fno-openmp-target-jit">, Group<f_Group>,
+  Flags<[CoreOption, NoArgumentUnused, HelpHidden]>;
 def fopenmp_target_new_runtime : Flag<["-"], "fopenmp-target-new-runtime">,
   Group<f_Group>, Flags<[CC1Option, HelpHidden]>;
 def fno_openmp_target_new_runtime : Flag<["-"], "fno-openmp-target-new-runtime">,
   Group<f_Group>, Flags<[CC1Option, HelpHidden]>;
 defm openmp_optimistic_collapse : BoolFOption<"openmp-optimistic-collapse",
   LangOpts<"OpenMPOptimisticCollapse">, DefaultFalse,
   PosFlag<SetTrue, [CC1Option]>, NegFlag<SetFalse>, BothFlags<[NoArgumentUnused, HelpHidden]>>;
 def static_openmp: Flag<["-"], "static-openmp">,

clang/lib/Driver/Driver.cpp

 
 // Parse the LTO options.
 void Driver::setLTOMode(const llvm::opt::ArgList &Args) {
   LTOMode =
       parseLTOMode(*this, Args, options::OPT_flto_EQ, options::OPT_fno_lto);
 
   OffloadLTOMode = parseLTOMode(*this, Args, options::OPT_foffload_lto_EQ,
                                 options::OPT_fno_offload_lto);
+
+  // Try to enable `-foffload-lto=full` if `-fopenmp-target-jit` is on.
+  if (Args.hasFlag(options::OPT_fopenmp_target_jit,
+                   options::OPT_fno_openmp_target_jit, false)) {
+    if (Arg *A = Args.getLastArg(options::OPT_foffload_lto_EQ,
+                                 options::OPT_fno_offload_lto))
+      if (OffloadLTOMode != LTOK_Full)
+        Diag(diag::err_drv_incompatible_options)
+            << A->getSpelling() << "-fopenmp-target-jit";
+    OffloadLTOMode = LTOK_Full;
+  }
 }
 
 /// Compute the desired OpenMP runtime from the flags provided.
 Driver::OpenMPRuntimeKind Driver::getOpenMPRuntime(const ArgList &Args) const {
   StringRef RuntimeName(CLANG_DEFAULT_OPENMP_RUNTIME);
 
   const Arg *A = Args.getLastArg(options::OPT_fopenmp_EQ);
   if (A)

clang/lib/Driver/ToolChains/Clang.cpp

for (Arg *A : Args.filtered(options::OPT_Xoffload_linker)) {
     else
       CmdArgs.push_back(Args.MakeArgString(
           "--device-linker=" +
           ToolChain::getOpenMPTriple(Val.drop_front()).getTriple() + "=" +
           A->getValue(1)));
   }
   Args.ClaimAllArgs(options::OPT_Xoffload_linker);
 
+  // Embed bitcode instead of an object in JIT mode.
+  if (const Arg *A = Args.getLastArg(options::OPT_fopenmp_target_jit))
+    CmdArgs.push_back("--embed-bitcode");
+
   // Forward `-mllvm` arguments to the LLVM invocations if present.
   for (Arg *A : Args.filtered(options::OPT_mllvm)) {
     CmdArgs.push_back("-mllvm");
     CmdArgs.push_back(A->getValue());
     A->claim();
   }
 
   // Add the linker arguments to be forwarded by the wrapper.

clang/test/Driver/openmp-offload-jit.c

+// REQUIRES: x86-registered-target
+// REQUIRES: nvptx-registered-target
+// REQUIRES: amdgpu-registered-target
+
+// Check that we enable LTO-mode properly with '-fopenmp-target-jit' and that it
+// still enabled LTO-mode if `-fno-offload-lto` is on.
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -ccc-print-phases -fopenmp \
+// RUN:   -fopenmp-targets=nvptx64-nvidia-cuda -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=PHASES-JIT %s
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -ccc-print-phases -fopenmp \
+// RUN:   -fopenmp-targets=nvptx64-nvidia-cuda -foffload-lto -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=PHASES-JIT %s
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -ccc-print-phases -fopenmp \
+// RUN:   -fopenmp-targets=amdgcn-amd-amdhsa -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=PHASES-JIT %s
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -ccc-print-phases -fopenmp \
+// RUN:   -fopenmp-targets=amdgcn-amd-amdhsa -foffload-lto -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=PHASES-JIT %s
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -ccc-print-phases -fopenmp \
+// RUN:   -fopenmp-targets=amdgcn-amd-amdhsa -fno-offload-lto -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=PHASES-JIT %s
+//
+//      PHASES-JIT: 0: input, "[[INPUT:.+]]", c, (host-openmp)
+// PHASES-JIT-NEXT: 1: preprocessor, {0}, cpp-output, (host-openmp)
+// PHASES-JIT-NEXT: 2: compiler, {1}, ir, (host-openmp)
+// PHASES-JIT-NEXT: 3: input, "[[INPUT]]", c, (device-openmp)
+// PHASES-JIT-NEXT: 4: preprocessor, {3}, cpp-output, (device-openmp)
+// PHASES-JIT-NEXT: 5: compiler, {4}, ir, (device-openmp)
+// PHASES-JIT-NEXT: 6: offload, "host-openmp (x86_64-unknown-linux-gnu)" {2}, "device-openmp ([[TARGET:.+]])" {5}, ir
+// PHASES-JIT-NEXT: 7: backend, {6}, lto-bc, (device-openmp)
+// PHASES-JIT-NEXT: 8: offload, "device-openmp ([[TARGET]])" {7}, lto-bc
+// PHASES-JIT-NEXT: 9: clang-offload-packager, {8}, image, (device-openmp)
+// PHASES-JIT-NEXT: 10: offload, "host-openmp (x86_64-unknown-linux-gnu)" {2}, "device-openmp (x86_64-unknown-linux-gnu)" {9}, ir
+// PHASES-JIT-NEXT: 11: backend, {10}, assembler, (host-openmp)
+// PHASES-JIT-NEXT: 12: assembler, {11}, object, (host-openmp)
+// PHASES-JIT-NEXT: 13: clang-linker-wrapper, {12}, image, (host-openmp)
+
+// Check that we add the `--embed-bitcode` flag to the linker wrapper.
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -fopenmp \
+// RUN:   -fopenmp-targets=nvptx64-nvidia-cuda -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=LINKER %s
+// LINKER: clang-linker-wrapper"{{.*}}"--embed-bitcode"
+
+// Check for incompatible combinations
+
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -fopenmp -fno-offload-lto \
+// RUN:   -fopenmp-targets=nvptx64-nvidia-cuda -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=NO-LTO %s
+// NO-LTO: error: The combination of '-fno-offload-lto' and '-fopenmp-target-jit' is incompatible
+
+// RUN: %clang -### --target=x86_64-unknown-linux-gnu -fopenmp -foffload-lto=thin \
+// RUN:   -fopenmp-targets=nvptx64-nvidia-cuda -fopenmp-target-jit %s 2>&1 \
+// RUN: | FileCheck -check-prefix=THIN-LTO %s
+// THIN-LTO: error: The combination of '-foffload-lto=' and '-fopenmp-target-jit' is incompatible