Allow aot_compiler_main to avoid compiling anything if only the protobuf descriptor is requested.

This lets us speed up compilation somewhat and enable LSPs to build jit wrappers without having to fully LLVM optimize/codegen the module.

PiperOrigin-RevId: 860324414

Author: allight

xls/build_rules/xls_internal_aot_rules.bzl

@@ -102,14 +102,31 @@ def _xls_aot_generate_impl(ctx):
     proto_file = ctx.actions.declare_file(out_proto_filename)
     obj_file = ctx.actions.declare_file(out_obj_filename)
     args = ctx.actions.args()
-    args.add("-input", src.ir_file.path)
+    skeleton_args = ctx.actions.args()
+
+    def common_add(*va, **kwargs):
+        args.add(*va, **kwargs)
+        skeleton_args.add(*va, **kwargs)
+
+    common_add("-input", src.ir_file.path)
     if (ctx.attr.salt_symbols):
-        args.add("-symbol_salt", str(ctx.label))
-    args.add("-top", ctx.attr.top)
+        common_add("-symbol_salt", str(ctx.label))
+    common_add("-aot_target", ctx.attr.aot_target)
+    common_add("-top", ctx.attr.top)
+    other_linking_contexts = []
+    if ctx.attr.with_msan:
+        common_add("--include_msan=true")
+
+        # With msan we need the TLS implementation.
+        other_linking_contexts = [ctx.attr._jit_emulated_tls[CcInfo].linking_context]
+    else:
+        common_add("--include_msan=false")
+
     args.add("-output_object", obj_file.path)
-    args.add("-output_proto", proto_file.path)
     args.add("-llvm_opt_level", ctx.attr.llvm_opt_level)
-    args.add("--aot_target", ctx.attr.aot_target)
+
+    skeleton_args.add("-output_proto", proto_file.path)
+
     extra_files = []
     aot_direct_request = ctx.attr._emit_aot_intermediates[BuildSettingInfo].value
     save_temps_reqest = ctx.attr._save_temps_is_requested[BoolConfigSettingInfo].value
@@ -124,23 +141,28 @@ def _xls_aot_generate_impl(ctx):
         args.add("-output_llvm_opt_ir", llvm_opt_ir_file.path)
         args.add("-output_asm", asm_file.path)
 
-    other_linking_contexts = []
-    if ctx.attr.with_msan:
-        args.add("--include_msan=true")
-
-        # With msan we need the TLS implementation.
-        other_linking_contexts = [ctx.attr._jit_emulated_tls[CcInfo].linking_context]
-    else:
-        args.add("--include_msan=false")
+    # Non-skeleton run to create the object file.
     ctx.actions.run(
-        outputs = [proto_file, obj_file] + extra_files,
+        outputs = [obj_file] + extra_files,
         inputs = [src.ir_file],
         arguments = [args],
         executable = aot_compiler,
         mnemonic = "AOTCompiling",
         progress_message = "Aot(JIT)Compiling %{label}: %{input}",
         toolchain = None,
     )
+
+    # Skeleton run to create the proto file.
+    ctx.actions.run(
+        outputs = [proto_file],
+        inputs = [src.ir_file],
+        arguments = [skeleton_args],
+        executable = aot_compiler,
+        mnemonic = "AotSkeleton",
+        progress_message = "Generating AOT skeleton %{label}: %{input}",
+        toolchain = None,
+    )
+
     obj_file_outputs = cc_common.create_compilation_outputs(
         objects = depset([obj_file]),
         pic_objects = depset([obj_file]),

xls/jit/BUILD

@@ -122,6 +122,7 @@ cc_binary(
         "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "//xls/dev_tools:extract_interface",
+        "//xls/interpreter:evaluator_options",
         "//xls/ir",
         "//xls/ir:block_elaboration",
         "//xls/ir:ir_parser",
@@ -748,6 +749,7 @@ cc_library(
         ":jit_evaluator_options",
         ":llvm_compiler",
         ":observer",
+        "//xls/common/status:ret_check",
         "//xls/common/status:status_macros",
         "@com_google_absl//absl/flags:flag",
         "@com_google_absl//absl/status",

xls/jit/aot_compiler.cc

@@ -14,12 +14,15 @@
 
 #include "xls/jit/aot_compiler.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/flags/flag.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "llvm/include/llvm/ADT/SmallVector.h"
 #include "llvm/include/llvm/ADT/StringRef.h"
@@ -46,6 +49,7 @@
 #include "llvm/include/llvm/TargetParser/Triple.h"
 #include "llvm/include/llvm/TargetParser/X86TargetParser.h"
 #include "llvm/include/llvm/Transforms/Utils/Cloning.h"
+#include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/jit/jit_emulated_tls.h"
 #include "xls/jit/jit_evaluator_options.h"
@@ -201,6 +205,16 @@ absl::Status AddWeakEmuTls(llvm::Module& module, llvm::LLVMContext* context) {
 absl::Status AotCompiler::CompileModule(
     std::unique_ptr<llvm::Module>&& module) {
   JitObserverRequests notification;
+  if (jit_options_.generate_skeleton()) {
+    // No need to actually compile anything.
+    object_code_.emplace();
+    XLS_RET_CHECK(jit_options_.jit_observer() == nullptr ||
+                  !jit_options_.jit_observer()
+                       ->GetNotificationOptions()
+                       .has_any_requests())
+        << "skeleton not compatible with observers";
+    return absl::OkStatus();
+  }
   if (jit_options_.jit_observer() != nullptr) {
     notification = jit_options_.jit_observer()->GetNotificationOptions();
   }

xls/jit/aot_compiler.h

@@ -62,6 +62,8 @@ class AotCompiler final : public LlvmCompiler {
     return *object_code_;
   }
 
+  bool is_skeleton() const override { return jit_options_.generate_skeleton(); }
+
  protected:
   absl::Status InitInternal() override;
 

xls/jit/aot_compiler_main.cc

@@ -39,6 +39,7 @@
 #include "xls/common/status/ret_check.h"
 #include "xls/common/status/status_macros.h"
 #include "xls/dev_tools/extract_interface.h"
+#include "xls/interpreter/evaluator_options.h"
 #include "xls/ir/block_elaboration.h"
 #include "xls/ir/function.h"
 #include "xls/ir/function_base.h"
@@ -59,6 +60,18 @@
 #include "xls/jit/type_buffer_metadata.h"
 #include "xls/jit/type_layout.pb.h"
 
+static constexpr std::string_view kUsage = R"(
+aot_compiler <flags>
+
+Compile the given IR file to an object file and a proto file.
+
+If --output_proto/--output_textproto is given and all of --output_object,
+--output_llvm_ir, --output_llvm_opt_ir, and --output_asm are not given, the
+compilation will not be performed and only the proto will be generated. This is
+significantly faster than performing the compilation and can be used to quickly
+check the output of the compiler.
+)";
+
 ABSL_FLAG(std::string, input, "", "Path to the IR to compile.");
 ABSL_FLAG(std::string, symbol_salt, "",
           "Additional text to append to symbol names to ensure no collisions.");
@@ -264,11 +277,14 @@ absl::Status RealMain(const std::string& input_ir_path,
   }
 
   std::optional<JitObjectCode> object_code;
+  bool generate_skeleton = !output_object_path && !output_llvm_ir_path &&
+                           !output_llvm_opt_ir_path && !output_asm_path;
   JitEvaluatorOptions jit_opts;
   jit_opts.set_opt_level(llvm_opt_level)
       .set_jit_observer(&obs)
       .set_symbol_salt(absl::GetFlag(FLAGS_symbol_salt))
-      .set_include_msan(include_msan);
+      .set_include_msan(include_msan)
+      .set_generate_skeleton(generate_skeleton);
   if (f->IsFunction()) {
     XLS_ASSIGN_OR_RETURN(
         object_code, FunctionJit::CreateObjectCode(
@@ -288,33 +304,35 @@ absl::Status RealMain(const std::string& input_ir_path,
     XLS_ASSIGN_OR_RETURN(object_code,
                          BlockJit::CreateObjectCode(elab, jit_opts));
   }
-  AotPackageEntrypointsProto all_entrypoints;
   if (output_object_path) {
     XLS_RETURN_IF_ERROR(SetFileContents(
         *output_object_path, std::string(object_code->object_code.begin(),
                                          object_code->object_code.end())));
   }
 
-  *all_entrypoints.mutable_data_layout() =
-      object_code->data_layout.getStringRepresentation();
+  if (output_proto_path || output_textproto_path) {
+    AotPackageEntrypointsProto all_entrypoints;
+    *all_entrypoints.mutable_data_layout() =
+        object_code->data_layout.getStringRepresentation();
 
-  auto context = std::make_unique<llvm::LLVMContext>();
-  LlvmTypeConverter type_converter(context.get(), object_code->data_layout);
-  for (const FunctionEntrypoint& oc : object_code->entrypoints) {
-    XLS_ASSIGN_OR_RETURN(
-        *all_entrypoints.add_entrypoint(),
-        GenerateEntrypointProto(
-            object_code->package ? object_code->package.get() : package.get(),
-            oc, include_msan, type_converter));
-  }
-  if (output_textproto_path) {
-    std::string text;
-    XLS_RET_CHECK(google::protobuf::TextFormat::PrintToString(all_entrypoints, &text));
-    XLS_RETURN_IF_ERROR(SetFileContents(*output_textproto_path, text));
-  }
-  if (output_proto_path) {
-    XLS_RETURN_IF_ERROR(SetFileContents(*output_proto_path,
-                                        all_entrypoints.SerializeAsString()));
+    auto context = std::make_unique<llvm::LLVMContext>();
+    LlvmTypeConverter type_converter(context.get(), object_code->data_layout);
+    for (const FunctionEntrypoint& oc : object_code->entrypoints) {
+      XLS_ASSIGN_OR_RETURN(
+          *all_entrypoints.add_entrypoint(),
+          GenerateEntrypointProto(
+              object_code->package ? object_code->package.get() : package.get(),
+              oc, include_msan, type_converter));
+    }
+    if (output_textproto_path) {
+      std::string text;
+      XLS_RET_CHECK(google::protobuf::TextFormat::PrintToString(all_entrypoints, &text));
+      XLS_RETURN_IF_ERROR(SetFileContents(*output_textproto_path, text));
+    }
+    if (output_proto_path) {
+      XLS_RETURN_IF_ERROR(SetFileContents(*output_proto_path,
+                                          all_entrypoints.SerializeAsString()));
+    }
   }
   if (output_llvm_ir_path) {
     XLS_RETURN_IF_ERROR(
@@ -335,7 +353,7 @@ absl::Status RealMain(const std::string& input_ir_path,
 }  // namespace xls
 
 int main(int argc, char** argv) {
-  xls::InitXls(argv[0], argc, argv);
+  xls::InitXls(kUsage, argc, argv);
   std::string input_ir_path = absl::GetFlag(FLAGS_input);
   QCHECK(!input_ir_path.empty()) << "--input must be specified.";
 

xls/jit/function_base_jit.cc

@@ -853,6 +853,15 @@ absl::StatusOr<PartitionedFunction> BuildFunctionInternal(
           .name = "continuation_point",
           .type = llvm::Type::getInt64Ty(jit_context.context())});
 
+  // NB Ideally we would be able to just immediately return here and not
+  // generate any of the llvm ir code or partition functions. Unfortunately we
+  // need to do that to figure out how big some of the buffers need to be. We
+  // still skip actually doing anything with the llvm bitcode we generate so the
+  // skeleton is still significantly faster than a full aot compile.
+  // TODO(https://github.com/google/xls/issues/3724): This shouldn't be needed.
+  // Ideally we should be able to determine the buffer sizes without writing out
+  // the llvm ir and return immediately.
+
   XLS_RETURN_IF_ERROR(AllocateBuffers(partitions, wrapper, allocator));
 
   std::vector<llvm::Function*> partition_functions;
@@ -1174,6 +1183,15 @@ absl::StatusOr<llvm::Function*> BuildPackedWrapper(
       LlvmFunctionWrapper::FunctionArg{
           .name = "continuation_point",
           .type = llvm::Type::getInt64Ty(*context)});
+  if (jit_context.is_skeleton()) {
+    // No need to actually generate any code for a skeleton compile
+    // Add a ret to ensure the code is well formed (though since we don't
+    // actually try to do anything with it this isn't necessary it still keeps
+    // it cleaner).
+    wrapper.entry_builder().CreateRet(
+        llvm::ConstantInt::get(llvm::Type::getInt64Ty(*context), /*value=*/0));
+    return wrapper.function();
+  }
 
   // First load and unpack the arguments then store them in LLVM native data
   // layout. These unpacked values are pointed to by an array of pointers passed
@@ -1333,6 +1351,8 @@ absl::StatusOr<JittedFunctionBase> JittedFunctionBase::BuildInternal(
 
   jitted_function.function_name_ = function_name;
   if (jit_context.llvm_compiler().IsOrcJit()) {
+    XLS_RET_CHECK(!jit_context.is_skeleton())
+        << "Only AOT can generate skeleton compiles";
     XLS_ASSIGN_OR_RETURN(auto* orc_jit, jit_context.llvm_compiler().AsOrcJit());
     XLS_ASSIGN_OR_RETURN(auto fn_address, orc_jit->LoadSymbol(function_name));
     jitted_function.function_ = absl::bit_cast<JitFunctionType>(fn_address);
@@ -1345,6 +1365,8 @@ absl::StatusOr<JittedFunctionBase> JittedFunctionBase::BuildInternal(
   if (build_packed_wrapper) {
     jitted_function.packed_function_name_ = packed_wrapper_name;
     if (jit_context.llvm_compiler().IsOrcJit()) {
+      XLS_RET_CHECK(!jit_context.is_skeleton())
+          << "Only AOT can generate skeleton compiles";
       XLS_ASSIGN_OR_RETURN(auto* orc_jit,
                            jit_context.llvm_compiler().AsOrcJit());
       XLS_ASSIGN_OR_RETURN(auto packed_fn_address,

xls/jit/ir_builder_visitor.h

@@ -61,6 +61,7 @@ class JitBuilderContext {
   LlvmCompiler& llvm_compiler() { return llvm_compiler_; }
   LlvmTypeConverter& type_converter() { return type_converter_; }
   FunctionBase* top() const { return top_; }
+  bool is_skeleton() const { return llvm_compiler_.is_skeleton(); }
 
   // Destructively returns the underlying llvm::Module.
   std::unique_ptr<llvm::Module> ConsumeModule() { return std::move(module_); }

xls/jit/jit_evaluator_options.h

@@ -66,12 +66,22 @@ class JitEvaluatorOptions {
   }
   JitObserver* jit_observer() const { return jit_observer_; }
 
+  // Tell the JIT to generate skeleton object code. Don't actually compile
+  // anything but just create the symbols we would generate if we were doing a
+  // full compile.
+  JitEvaluatorOptions& set_generate_skeleton(bool value) {
+    generate_skeleton_ = value;
+    return *this;
+  }
+  bool generate_skeleton() const { return generate_skeleton_; }
+
  private:
   int64_t opt_level_ = LlvmCompiler::kDefaultOptLevel;
   std::string symbol_salt_;
   bool include_observer_callbacks_ = false;
   bool include_msan_ = false;
   JitObserver* jit_observer_ = nullptr;
+  bool generate_skeleton_ = false;
 };
 
 }  // namespace xls

xls/jit/llvm_compiler.h

@@ -92,6 +92,10 @@ class LlvmCompiler {
     return include_observer_callbacks_;
   }
 
+  // Return true if this is a skeleton compilation. That is don't actually
+  // compile anything just create the symbols.
+  virtual bool is_skeleton() const { return false; }
+
  protected:
   absl::Status Init();
 

xls/jit/observer.h

@@ -45,6 +45,10 @@ struct JitObserverRequests {
   bool optimized_module = false;
   // Do we want to get called with optimized asm code.
   bool assembly_code_str = false;
+
+  bool has_any_requests() const {
+    return unoptimized_module || optimized_module || assembly_code_str;
+  }
 };
 
 // Basic observer for JIT compilation events