[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

1 files changed, 22 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
index fbc071536d22..6b1b9520972e 100644
--- a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp
+++ b/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 @@ public:
   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.