[CIR][MERGE-SPLIT-COMPILATION] Adds a compilation step that merges Device and Host CIR into a single module and can co-optimize their execution

clang/include/clang/CIR/FrontendAction/CIRGenAction.h

@@ -113,6 +113,23 @@ class EmitObjAction : public CIRGenAction {
   EmitObjAction(mlir::MLIRContext *mlirCtx = nullptr);
 };
 
+class CIRCombineAction : public clang::FrontendAction {
+public:
+  std::unique_ptr<clang::ASTConsumer>
+  CreateASTConsumer(clang::CompilerInstance &CI,
+                    llvm::StringRef InFile) override {
+    return std::make_unique<clang::ASTConsumer>();
+  }
+
+  void ExecuteAction() override {
+    llvm::report_fatal_error("CIRCombineAction NIY");
+  }
+
+  // We don't need a preprocessor-only mode.
+  bool usesPreprocessorOnly() const override { return false; }
+  virtual bool hasCIRSupport() const override { return true; }
+};
+
 // Used for -fclangir-analysis-only: use CIR analysis but still use original
 // LLVM codegen path
 class AnalysisOnlyActionBase : public clang::CodeGenAction {

clang/include/clang/Driver/Options.td

@@ -3208,6 +3208,7 @@ def fclangir_mem2reg : Flag<["-"], "fclangir-mem2reg">,
   HelpText<"Enable mem2reg on the flat ClangIR">,
   MarshallingInfoFlag<FrontendOpts<"ClangIREnableMem2Reg">>;
 
+def CIR_Group : OptionGroup<"<CIR options>">;
 def clangir_disable_passes : Flag<["-"], "clangir-disable-passes">,
   Visibility<[ClangOption, CC1Option]>,
   HelpText<"Disable CIR transformations pipeline">,
@@ -3256,6 +3257,16 @@ def emit_cir_only : Flag<["-"], "emit-cir-only">,
 def emit_cir_flat : Flag<["-"], "emit-cir-flat">, Visibility<[ClangOption, CC1Option]>,
   Group<Action_Group>, Alias<emit_mlir_EQ>, AliasArgs<["cir-flat"]>,
   HelpText<"Similar to -emit-cir but also lowers structured CFG into basic blocks.">;
+def cir_combine : Flag<["-"], "cir-combine">,
+  Visibility<[CC1Option]>,
+  Group<Action_Group>,
+  HelpText<"Combine host and device CIR modules into a single offload container CIR module.">;
+def cir_host_input : Separate<["-"], "cir-host-input">,
+  Visibility<[CC1Option]>, Group<CIR_Group>,
+  HelpText<"Host CIR input for -cir-combine.">;
+def cir_device_input : Separate<["-"], "cir-device-input">,
+  Visibility<[CC1Option]>, Group<CIR_Group>,
+  HelpText<"Device CIR input for -cir-combine (may be repeated).">;
 /// ClangIR-specific options - END
 
 def flto : Flag<["-"], "flto">,

clang/include/clang/Frontend/FrontendOptions.h

@@ -68,6 +68,10 @@ enum ActionKind {
   /// Generate CIR, bud don't emit anything.
   EmitCIROnly,
 
+  /// Combine multiple CIR modules (e.g. host and device) into a single
+  /// container
+  CIRCombine,
+
   /// Emit a .mlir file
   EmitMLIR,
 
@@ -454,6 +458,8 @@ class FrontendOptions {
   std::string ClangIRIdiomRecognizerOpts;
   std::string ClangIRLibOptOpts;
   std::string ClangIRFile;
+  std::string CIRHostInput;
+  std::string CIRDeviceInput;
 
   frontend::MLIRDialectKind MLIRTargetDialect = frontend::MLIR_CORE;
 

clang/lib/Frontend/CompilerInvocation.cpp

@@ -2790,6 +2790,7 @@ static const auto &getFrontendActionTable() {
       {frontend::EmitCodeGenOnly, OPT_emit_codegen_only},
       {frontend::EmitObj, OPT_emit_obj},
       {frontend::ExtractAPI, OPT_extract_api},
+      {frontend::CIRCombine, OPT_cir_combine},
 
       {frontend::FixIt, OPT_fixit_EQ},
       {frontend::FixIt, OPT_fixit},
@@ -2901,6 +2902,13 @@ static void GenerateFrontendArgs(const FrontendOptions &Opts,
     };
   }
 
+  if (Opts.ProgramAction == frontend::CIRCombine) {
+    if (!Opts.CIRHostInput.empty())
+      GenerateArg(Consumer, OPT_cir_host_input, Opts.CIRHostInput);
+    if (!Opts.CIRDeviceInput.empty())
+      GenerateArg(Consumer, OPT_cir_device_input, Opts.CIRDeviceInput);
+  }
+
   if (Opts.ProgramAction == frontend::FixIt && !Opts.FixItSuffix.empty()) {
     GenerateProgramAction = [&]() {
       GenerateArg(Consumer, OPT_fixit_EQ, Opts.FixItSuffix);
@@ -3001,6 +3009,8 @@ static void GenerateFrontendArgs(const FrontendOptions &Opts,
       Lang = "assembler-with-cpp";
       break;
     case Language::Unknown:
+      if (Opts.ProgramAction == frontend::CIRCombine)
+        break;
       assert(Opts.DashX.getFormat() == InputKind::Precompiled &&
              "Generating -x argument for unknown language (not precompiled).");
       Lang = "ast";
@@ -3016,8 +3026,9 @@ static void GenerateFrontendArgs(const FrontendOptions &Opts,
       break;
     }
 
-    GenerateArg(Consumer, OPT_x,
-                Lang + HeaderUnit + Header + ModuleMap + Preprocessed);
+    if (Opts.ProgramAction != frontend::CIRCombine)
+      GenerateArg(Consumer, OPT_x,
+                  Lang + HeaderUnit + Header + ModuleMap + Preprocessed);
   }
 
   // OPT_INPUT has a unique class, generate it directly.
@@ -3173,6 +3184,37 @@ static bool ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
   if (Args.hasArg(OPT_fclangir_direct_lowering))
     Opts.ClangIRDirectLowering = true;
 
+  if (Args.hasArg(OPT_cir_combine) && !Args.hasArg(OPT_fclangir))
+    Diags.Report(diag::err_drv_argument_only_allowed_with)
+        << "-cir-combine" << "-fclangir";
+
+  if (Args.hasArg(OPT_cir_combine) &&
+      Opts.ProgramAction == frontend::CIRCombine) {
+    if (!Opts.Inputs.empty()) {
+      Diags.Report(diag::err_drv_invalid_argument_to_option)
+          << "positional inputs" << "-cir-combine";
+    }
+
+    auto HostArgs = Args.filtered(OPT_cir_host_input);
+    unsigned HostCount = std::distance(HostArgs.begin(), HostArgs.end());
+    if (HostCount == 0)
+      Diags.Report(diag::err_drv_missing_argument) << "-cir-host-input" << 1;
+    if (HostCount > 1)
+      Diags.Report(diag::err_drv_invalid_argument_to_option)
+          << "multiple -cir-host-input" << "-cir-combine";
+    Opts.CIRHostInput = Args.getLastArgValue(OPT_cir_host_input).str();
+
+    auto DevArgs = Args.filtered(OPT_cir_device_input);
+    unsigned DevCount = std::distance(DevArgs.begin(), DevArgs.end());
+    if (DevCount == 0)
+      Diags.Report(diag::err_drv_missing_argument) << "-cir-device-input" << 1;
+
+    if (DevCount > 1)
+      Diags.Report(diag::err_drv_invalid_argument_to_option)
+          << "multiple -cir-device-input" << "-cir-combine";
+    Opts.CIRDeviceInput = Args.getLastArgValue(OPT_cir_device_input).str();
+  }
+
   if (Args.hasArg(OPT_clangir_disable_passes))
     Opts.ClangIRDisablePasses = true;
 
@@ -3301,6 +3343,14 @@ static bool ParseFrontendArgs(FrontendOptions &Opts, ArgList &Args,
   // '-' is the default input if none is given.
   std::vector<std::string> Inputs = Args.getAllArgValues(OPT_INPUT);
   Opts.Inputs.clear();
+  if (Opts.ProgramAction == frontend::CIRCombine) {
+    Opts.Inputs.emplace_back(
+        Opts.CIRHostInput,
+        clang::InputKind(clang::Language::CIR, clang::InputKind::Source),
+        false);
+    // CIRCombine does not require the -x flag for now. We skip this part
+    return Diags.getNumErrors() == NumErrorsBefore;
+  }
   if (Inputs.empty())
     Inputs.push_back("-");
 
@@ -4721,6 +4771,7 @@ static bool isStrictlyPreprocessorAction(frontend::ActionKind Action) {
   case frontend::ASTDump:
   case frontend::ASTPrint:
   case frontend::ASTView:
+  case frontend::CIRCombine:
   case frontend::EmitAssembly:
   case frontend::EmitBC:
   case frontend::EmitCIROnly:

clang/lib/FrontendTool/ExecuteCompilerInvocation.cpp

@@ -73,8 +73,16 @@ CreateFrontendBaseAction(CompilerInstance &CI) {
   case ASTView:                return std::make_unique<ASTViewAction>();
   case DumpCompilerOptions:
     return std::make_unique<DumpCompilerOptionsAction>();
-  case DumpRawTokens:          return std::make_unique<DumpRawTokensAction>();
-  case DumpTokens:             return std::make_unique<DumpTokensAction>();
+  case DumpRawTokens:
+    return std::make_unique<DumpRawTokensAction>();
+  case DumpTokens:
+    return std::make_unique<DumpTokensAction>();
+  case CIRCombine:
+#if CLANG_ENABLE_CIR
+    return std::make_unique<cir::CIRCombineAction>();
+#else
+    llvm_unreachable("CIR suppport not built into clang");
+#endif
   case EmitAssembly:
 #if CLANG_ENABLE_CIR
     if (UseCIR)

clang/tools/cir-host-device-bundle/CMakeLists.txt

@@ -0,0 +1,45 @@
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+
+include_directories(${LLVM_MAIN_SRC_DIR}/../mlir/include)
+include_directories(${CMAKE_BINARY_DIR}/tools/mlir/include)
+
+# GCC, unlike clang, issues a warning when one virtual function is overridden
+# in a derived class but one or more other virtual functions with the same
+# name and different signature from a base class are not overridden. This
+# leads to many warnings in the MLIR and ClangIR code when using the
+# OpenConversionPattern<>::matchAndRewrite() function in the ordinary way.
+# The "hiding" behavior is what we want, so we're just disabling the warning
+# here.
+if (LLVM_COMPILER_IS_GCC_COMPATIBLE AND (NOT "${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang"))
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-overloaded-virtual")
+endif()
+
+add_clang_tool(cir-host-device-bundle
+  cir-host-device-bundle.cpp
+)
+
+clang_target_link_libraries(cir-host-device-bundle
+  PRIVATE
+  clangCIR
+  clangCIRLoweringDirectToLLVM
+  MLIRCIR
+  MLIRCIRTransforms
+)
+
+target_link_libraries(cir-host-device-bundle
+  PRIVATE
+  ${dialect_libs}
+  ${conversion_libs}
+  MLIRAnalysis
+  MLIRDialect
+  MLIRIR
+  MLIRMemRefDialect
+  MLIROptLib
+  MLIRParser
+  MLIRPass
+  MLIRRegisterAllPasses
+  MLIRSideEffectInterfaces
+  MLIRTransforms
+  MLIRTransformUtils
+)