Add --llvm_opt_level option to aot_compiler_main

This can be useful if the llvm optimizer is crashing to aid in debugging.

Bug: #1776
PiperOrigin-RevId: 707988660

Author: allight

xls/jit/BUILD

@@ -99,6 +99,7 @@ cc_binary(
         ":function_base_jit",
         ":function_jit",
         ":jit_proc_runtime",
+        ":llvm_compiler",
         ":llvm_type_converter",
         ":observer",
         ":type_layout_cc_proto",

xls/jit/aot_compiler_main.cc

@@ -52,6 +52,7 @@
 #include "xls/jit/function_base_jit.h"
 #include "xls/jit/function_jit.h"
 #include "xls/jit/jit_proc_runtime.h"
+#include "xls/jit/llvm_compiler.h"
 #include "xls/jit/llvm_type_converter.h"
 #include "xls/jit/observer.h"
 #include "xls/jit/type_layout.pb.h"
@@ -75,6 +76,9 @@ ABSL_FLAG(std::optional<std::string>, output_llvm_opt_ir, std::nullopt,
           "Path at which to write the output optimized llvm file.");
 ABSL_FLAG(std::optional<std::string>, output_asm, std::nullopt,
           "Path at which to write the output optimized llvm file.");
+ABSL_FLAG(int64_t, llvm_opt_level, xls::LlvmCompiler::kDefaultOptLevel,
+          "The optimization level to use for the LLVM optimizer.");
+
 #ifdef ABSL_HAVE_MEMORY_SANITIZER
 static constexpr bool kHasMsan = true;
 #else
@@ -227,7 +231,7 @@ absl::Status RealMain(const std::string& input_ir_path,
                       const std::optional<std::string>& top,
                       const std::optional<std::string>& output_object_path,
                       const std::optional<std::string>& output_proto_path,
-                      bool include_msan,
+                      bool include_msan, int64_t llvm_opt_level,
                       const std::optional<std::string>& output_textproto_path,
                       const std::optional<std::string>& output_llvm_ir_path,
                       const std::optional<std::string>& output_llvm_opt_ir_path,
@@ -258,26 +262,26 @@ absl::Status RealMain(const std::string& input_ir_path,
 
   std::optional<JitObjectCode> object_code;
   if (f->IsFunction()) {
-    XLS_ASSIGN_OR_RETURN(
-        object_code,
-        FunctionJit::CreateObjectCode(f->AsFunctionOrDie(),
-                                      /*opt_level = */ 3, include_msan, &obs));
+    XLS_ASSIGN_OR_RETURN(object_code, FunctionJit::CreateObjectCode(
+                                          f->AsFunctionOrDie(), llvm_opt_level,
+                                          include_msan, &obs));
   } else if (f->IsProc()) {
     if (f->AsProcOrDie()->is_new_style_proc()) {
       XLS_ASSIGN_OR_RETURN(
-          object_code,
-          CreateProcAotObjectCode(f->AsProcOrDie(), include_msan, &obs));
+          object_code, CreateProcAotObjectCode(f->AsProcOrDie(), llvm_opt_level,
+                                               include_msan, &obs));
     } else {
       // all procs
-      XLS_ASSIGN_OR_RETURN(object_code, CreateProcAotObjectCode(
-                                            package.get(), include_msan, &obs));
+      XLS_ASSIGN_OR_RETURN(
+          object_code, CreateProcAotObjectCode(package.get(), llvm_opt_level,
+                                               include_msan, &obs));
     }
   } else {
     XLS_ASSIGN_OR_RETURN(BlockElaboration elab,
                          BlockElaboration::Elaborate(f->AsBlockOrDie()));
     XLS_ASSIGN_OR_RETURN(
         object_code,
-        BlockJit::CreateObjectCode(elab, /*opt_level=*/3, include_msan, &obs));
+        BlockJit::CreateObjectCode(elab, llvm_opt_level, include_msan, &obs));
   }
   AotPackageEntrypointsProto all_entrypoints;
   if (output_object_path) {
@@ -340,6 +344,7 @@ int main(int argc, char** argv) {
   bool include_msan = absl::GetFlag(FLAGS_include_msan);
   absl::Status status = xls::RealMain(
       input_ir_path, top, output_object_path, output_proto_path, include_msan,
+      absl::GetFlag(FLAGS_llvm_opt_level),
       absl::GetFlag(FLAGS_output_textproto),
       absl::GetFlag(FLAGS_output_llvm_ir),
       absl::GetFlag(FLAGS_output_llvm_opt_ir), absl::GetFlag(FLAGS_output_asm));

xls/jit/jit_proc_runtime.cc

@@ -124,12 +124,11 @@ class SharedCompiler final : public LlvmCompiler {
 };
 
 absl::StatusOr<JitObjectCode> GetAotObjectCode(ProcElaboration elaboration,
+                                               int64_t opt_level,
                                                bool with_msan,
                                                JitObserver* observer) {
-  XLS_ASSIGN_OR_RETURN(
-      std::unique_ptr<AotCompiler> compiler,
-      AotCompiler::Create(
-          with_msan, /*opt_level=*/LlvmCompiler::kDefaultOptLevel, observer));
+  XLS_ASSIGN_OR_RETURN(std::unique_ptr<AotCompiler> compiler,
+                       AotCompiler::Create(with_msan, opt_level, observer));
   XLS_ASSIGN_OR_RETURN(std::unique_ptr<llvm::TargetMachine> target,
                        compiler->CreateTargetMachine());
   llvm::DataLayout layout = target->createDataLayout();
@@ -288,17 +287,22 @@ absl::StatusOr<std::unique_ptr<SerialProcRuntime>> CreateJitSerialProcRuntime(
 }
 
 absl::StatusOr<JitObjectCode> CreateProcAotObjectCode(Package* package,
+                                                      int64_t opt_level,
                                                       bool with_msan,
                                                       JitObserver* observer) {
   XLS_ASSIGN_OR_RETURN(ProcElaboration elaboration,
                        ProcElaboration::ElaborateOldStylePackage(package));
-  return GetAotObjectCode(std::move(elaboration), with_msan, observer);
+  return GetAotObjectCode(std::move(elaboration), opt_level, with_msan,
+                          observer);
 }
-absl::StatusOr<JitObjectCode> CreateProcAotObjectCode(Proc* top, bool with_msan,
+absl::StatusOr<JitObjectCode> CreateProcAotObjectCode(Proc* top,
+                                                      int64_t opt_level,
+                                                      bool with_msan,
                                                       JitObserver* observer) {
   XLS_ASSIGN_OR_RETURN(ProcElaboration elaboration,
                        ProcElaboration::Elaborate(top));
-  return GetAotObjectCode(std::move(elaboration), with_msan, observer);
+  return GetAotObjectCode(std::move(elaboration), opt_level, with_msan,
+                          observer);
 }
 
 // Create a SerialProcRuntime composed of ProcJits. Constructed from the

xls/jit/jit_proc_runtime.h

@@ -15,6 +15,7 @@
 #ifndef XLS_JIT_JIT_PROC_RUNTIME_H_
 #define XLS_JIT_JIT_PROC_RUNTIME_H_
 
+#include <cstdint>
 #include <memory>
 #include <optional>
 
@@ -71,10 +72,12 @@ absl::StatusOr<std::unique_ptr<SerialProcRuntime>> CreateAotSerialProcRuntime(
 
 // Generate AOT code for the given proc elaboration.
 absl::StatusOr<JitObjectCode> CreateProcAotObjectCode(
-    Package* package, bool with_msan, JitObserver* observer = nullptr);
+    Package* package, int64_t opt_level, bool with_msan,
+    JitObserver* observer = nullptr);
 // Generate AOT code for the given proc elaboration.
 absl::StatusOr<JitObjectCode> CreateProcAotObjectCode(
-    Proc* top, bool with_msan, JitObserver* observer = nullptr);
+    Proc* top, int64_t opt_level, bool with_msan,
+    JitObserver* observer = nullptr);
 
 }  // namespace xls