[Driver][CodeGen] Properly handle -fsplit-machine-functions for fatbinary compilation.

When building a fatbinary, the driver invokes the compiler multiple
times with different "--target". (For example, with "-x cuda
--cuda-gpu-arch=sm_70" flags, clang will be invoded twice, once with
--target=x86_64_...., once with --target=sm_70) If we use
-fsplit-machine-functions or -fno-split-machine-functions for such
invocation, the driver reports an error.

This CL changes the behavior so:

  - "-fsplit-machine-functions" is now passed to all targets, for non-X86
    targets, the flag is a NOOP and causes a warning.

  - "-fno-split-machine-functions" now negates -fsplit-machine-functions (if
    -fno-split-machine-functions appears after any -fsplit-machine-functions)
    for any target triple, previously, it causes an error.

  - "-fsplit-machine-functions -Xarch_device -fno-split-machine-functions"
    enables MFS on host but disables MFS for GPUS without warnings/errors.

  - "-Xarch_host -fsplit-machine-functions" enables MFS on host but disables
    MFS for GPUS without warnings/errors.

Reviewed by: xur, dhoekwater

Differential Revision: https://reviews.llvm.org/D157750

Author: shenhanc78

clang/include/clang/Basic/DiagnosticDriverKinds.td

@@ -693,6 +693,10 @@ def warn_drv_fjmc_for_elf_only : Warning<
   "-fjmc works only for ELF; option ignored">,
   InGroup<OptionIgnored>;
 
+def warn_drv_for_elf_only : Warning<
+  "'%0' works only for ELF; option ignored">,
+  InGroup<OptionIgnored>;
+
 def warn_target_override_arm64ec : Warning<
   "/arm64EC has been overridden by specified target: %0; option ignored">,
   InGroup<OptionIgnored>;

clang/lib/Driver/ToolChains/Clang.cpp

@@ -5870,14 +5870,15 @@ void Clang::ConstructJob(Compilation &C, const JobAction &JA,
 
   if (Arg *A = Args.getLastArg(options::OPT_fsplit_machine_functions,
                                options::OPT_fno_split_machine_functions)) {
-    // This codegen pass is only available on x86-elf targets.
-    if (Triple.isX86() && Triple.isOSBinFormatELF()) {
-      if (A->getOption().matches(options::OPT_fsplit_machine_functions))
+    if (A->getOption().matches(options::OPT_fsplit_machine_functions)) {
+      // This codegen pass is only available on elf targets.
+      if (Triple.isOSBinFormatELF())
         A->render(Args, CmdArgs);
-    } else {
-      D.Diag(diag::err_drv_unsupported_opt_for_target)
-          << A->getAsString(Args) << TripleStr;
+      else
+        D.Diag(diag::warn_drv_for_elf_only) << A->getAsString(Args);
     }
+    // Do not issue warnings for -fno-split-machine-functions even it is not
+    // on ELF.
   }
 
   Args.AddLastArg(CmdArgs, options::OPT_finstrument_functions,

clang/test/Driver/fsplit-machine-functions-with-cuda-nvptx.c

@@ -0,0 +1,61 @@
+// REQUIRES: system-linux
+// REQUIRES: x86-registered-target
+// REQUIRES: nvptx-registered-target
+// REQUIRES: shell
+
+// Check that -fsplit-machine-functions is passed to both x86 and cuda
+// compilation and does not cause driver error.
+// RUN:   %clang -### --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -fsplit-machine-functions -S %s 2>&1 \
+// RUN:     | FileCheck %s --check-prefix=MFS1
+// MFS1: "-target-cpu" "x86-64"{{.*}}"-fsplit-machine-functions"
+// MFS1: "-target-cpu" "sm_70"{{.*}}"-fsplit-machine-functions"
+
+// Check that -fsplit-machine-functions is passed to cuda and it
+// causes a warning.
+// RUN:   %clang --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -fsplit-machine-functions -S %s 2>&1 \
+// RUN:     | FileCheck %s --check-prefix=MFS2
+// MFS2: warning: -fsplit-machine-functions is not valid for nvptx
+
+// Check that -Xarch_host -fsplit-machine-functions is passed only to
+// native compilation.
+// RUN:   %clang -### --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -Xarch_host \
+// RUN:     -fsplit-machine-functions -S %s \
+// RUN:     2>&1 | FileCheck %s --check-prefix=MFS3
+// MFS3:     "-target-cpu" "x86-64"{{.*}}"-fsplit-machine-functions"
+// MFS3-NOT: "-target-cpu" "sm_70"{{.*}}"-fsplit-machine-functions"
+
+// Check that -Xarch_host -fsplit-machine-functions does not cause any warning.
+// RUN:   %clang --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN      --cuda-gpu-arch=sm_70 -x cuda -Xarch_host \
+// RUN      -fsplit-machine-functions -S %s || { echo \
+// RUN      "warning: -fsplit-machine-functions is not valid for" ; } \
+// RUN      2>&1 | FileCheck %s --check-prefix=MFS4
+// MFS4-NOT: warning: -fsplit-machine-functions is not valid for
+
+// Check that -Xarch_device -fsplit-machine-functions does cause the warning.
+// RUN:   %clang --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -Xarch_device \
+// RUN:     -fsplit-machine-functions -S %s 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=MFS5
+// MFS5: warning: -fsplit-machine-functions is not valid for
+
+// Check that -fsplit-machine-functions -Xarch_device
+// -fno-split-machine-functions only passes MFS to x86
+// RUN:   %clang -### --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -fsplit-machine-functions \
+// RUN:     -Xarch_device -fno-split-machine-functions -S %s \
+// RUN:     2>&1 | FileCheck %s --check-prefix=MFS6
+// MFS6:     "-target-cpu" "x86-64"{{.*}}"-fsplit-machine-functions"
+// MFS6-NOT: "-target-cpu" "sm_70"{{.*}}"-fsplit-machine-functions"
+
+// Check that -fsplit-machine-functions -Xarch_device
+// -fno-split-machine-functions has no warnings
+// RUN:   %clang --target=x86_64-unknown-linux-gnu -nogpulib -nogpuinc \
+// RUN:     --cuda-gpu-arch=sm_70 -x cuda -fsplit-machine-functions \
+// RUN:     -Xarch_device -fno-split-machine-functions -S %s \
+// RUN:     || { echo "warning: -fsplit-machine-functions is not valid for"; } \
+// RUN:     2>&1 | FileCheck %s --check-prefix=MFS7
+// MFS7-NOT: warning: -fsplit-machine-functions is not valid for

clang/test/Driver/fsplit-machine-functions.c

@@ -1,8 +1,8 @@
 // RUN: %clang -### -target x86_64 -fprofile-use=default.profdata -fsplit-machine-functions %s -c 2>&1 | FileCheck -check-prefix=CHECK-OPT %s
 // RUN: %clang -### -target x86_64 -fsplit-machine-functions %s -c 2>&1 | FileCheck -check-prefix=CHECK-OPT %s
 // RUN: %clang -### -target x86_64 -fprofile-use=default.profdata -fsplit-machine-functions -fno-split-machine-functions %s -c 2>&1 | FileCheck -check-prefix=CHECK-NOOPT %s
-// RUN: not %clang -c -target arm-unknown-linux -fsplit-machine-functions %s 2>&1 | FileCheck -check-prefix=CHECK-TRIPLE %s
+// RUN: %clang -c -target arm-unknown-linux -fsplit-machine-functions %s 2>&1 | FileCheck -check-prefix=CHECK-TRIPLE %s
 
 // CHECK-OPT:       "-fsplit-machine-functions"
 // CHECK-NOOPT-NOT: "-fsplit-machine-functions"
-// CHECK-TRIPLE:    error: unsupported option '-fsplit-machine-functions' for target
+// CHECK-TRIPLE:    warning: -fsplit-machine-functions is not valid for arm

clang/test/Driver/fsplit-machine-functions2.c

@@ -7,6 +7,9 @@
 // Test the mix of -fsplit-machine-functions and -fno-split-machine-functions
 // RUN: %clang -### -target x86_64-unknown-linux -flto -fsplit-machine-functions -fno-split-machine-functions %s 2>&1 | FileCheck %s -check-prefix=CHECK-NOPASS
 // RUN: %clang -### -target x86_64-unknown-linux -flto -fno-split-machine-functions -fsplit-machine-functions %s 2>&1 | FileCheck %s -check-prefix=CHECK-PASS
+// Check that for non-X86, passing no-split-machine-functions does not cause error.
+// RUN: %clang -### -target aarch64-unknown-linux -flto -fsplit-machine-functions -fno-split-machine-functions %s 2>&1 | FileCheck %s -check-prefix=CHECK-NOPASS2
 
 // CHECK-PASS:          "-plugin-opt=-split-machine-functions"
 // CHECK-NOPASS-NOT:    "-plugin-opt=-split-machine-functions"
+// CHECK-NOPASS2-NOT:   "-plugin-opt=-split-machine-functions"

llvm/include/llvm/IR/DiagnosticInfo.h

@@ -86,6 +86,7 @@ enum DiagnosticKind {
   DK_SrcMgr,
   DK_DontCall,
   DK_MisExpect,
+  DK_MachineFunctionSplit,
   DK_FirstPluginKind // Must be last value to work with
                      // getNextAvailablePluginDiagnosticKind
 };
@@ -1117,6 +1118,20 @@ class DiagnosticInfoDontCall : public DiagnosticInfo {
   }
 };
 
+class DiagnosticInfoMachineFunctionSplit : public DiagnosticInfo {
+  StringRef TargetTriple;
+
+public:
+  DiagnosticInfoMachineFunctionSplit(StringRef TargetTriple,
+                                     DiagnosticSeverity DS)
+      : DiagnosticInfo(DK_MachineFunctionSplit, DS),
+        TargetTriple(TargetTriple) {}
+  void print(DiagnosticPrinter &DP) const override;
+  static bool classof(const DiagnosticInfo *DI) {
+    return DI->getKind() == DK_MachineFunctionSplit;
+  }
+};
+
 } // end namespace llvm
 
 #endif // LLVM_IR_DIAGNOSTICINFO_H

llvm/lib/CodeGen/MachineFunctionSplitter.cpp

@@ -35,9 +35,11 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/TargetParser/Triple.h"
 #include <optional>
 
 using namespace llvm;
@@ -82,6 +84,13 @@ class MachineFunctionSplitter : public MachineFunctionPass {
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 
   bool runOnMachineFunction(MachineFunction &F) override;
+
+  bool doInitialization(Module &) override;
+
+  static bool isSupportedTriple(const Triple &T) { return T.isX86(); }
+
+private:
+  bool UnsupportedTriple = false;
 };
 } // end anonymous namespace
 
@@ -127,7 +136,20 @@ static bool isColdBlock(const MachineBasicBlock &MBB,
   return (*Count < ColdCountThreshold);
 }
 
+bool MachineFunctionSplitter::doInitialization(Module &M) {
+  StringRef T = M.getTargetTriple();
+  if (!isSupportedTriple(Triple(T))) {
+    UnsupportedTriple = true;
+    M.getContext().diagnose(
+        DiagnosticInfoMachineFunctionSplit(T, DS_Warning));
+    return false;
+  }
+  return MachineFunctionPass::doInitialization(M);
+}
+
 bool MachineFunctionSplitter::runOnMachineFunction(MachineFunction &MF) {
+  if (UnsupportedTriple)
+    return false;
   // We target functions with profile data. Static information in the form
   // of exception handling code may be split to cold if user passes the
   // mfs-split-ehcode flag.

llvm/lib/IR/DiagnosticInfo.cpp

@@ -449,3 +449,7 @@ void DiagnosticInfoDontCall::print(DiagnosticPrinter &DP) const {
   if (!getNote().empty())
     DP << ": " << getNote();
 }
+
+void DiagnosticInfoMachineFunctionSplit::print(DiagnosticPrinter &DP) const {
+  DP << "-fsplit-machine-functions is not valid for " << TargetTriple;
+}

llvm/test/CodeGen/Generic/Inputs/fsloader-mfs.afdo

@@ -10,6 +10,15 @@
 ; RUN: sed 's/InstrProf/SampleProfile/g' %s > %t.ll
 ; RUN: llc < %t.ll -mtriple=x86_64-unknown-linux-gnu -split-machine-functions | FileCheck %s --check-prefix=FSAFDO-MFS
 ; RUN: llc < %t.ll -mtriple=x86_64-unknown-linux-gnu -split-machine-functions | FileCheck %s --check-prefix=FSAFDO-MFS2
+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu -debug-pass=Structure -fs-profile-file=%S/Inputs/fsloader-mfs.afdo -enable-fs-discriminator=true -improved-fs-discriminator=true -split-machine-functions 2>&1 | FileCheck %s --check-prefix=MFS_ON
+; RUN: llc < %s -mtriple=aarch64-unknown-linux-gnu -debug-pass=Structure -fs-profile-file=%S/Inputs/fsloader-mfs.afdo -enable-fs-discriminator=true -improved-fs-discriminator=true -split-machine-functions 2>&1 | FileCheck %s --check-prefix=MFS_OFF
+
+;; Check that MFS is on for X86 targets.
+; MFS_ON: Machine Function Splitter Transformation
+; MFS_ON_NO: warning: -fsplit-machine-functions is not valid for
+;; Check that MFS is not on for non-X86 targets.
+; MFS_OFF: warning: -fsplit-machine-functions is not valid for
+; MFS_OFF_NO: Machine Function Splitter Transformation
 
 define void @foo1(i1 zeroext %0) nounwind !prof !14 !section_prefix !15 {
 ;; Check that cold block is moved to .text.split.