commit 9138128fe683da5ab133d9baaa297d3eb110a3b7
author Stella Laurenzo <stellaraccident@gmail.com> Mon Jul 31 10:56:06 2023 -0700
committer GitHub <noreply@github.com> Mon Jul 31 10:56:06 2023 -0700
tree d3a9adbc33f532f6a2591d43d7cefca6d29aa27e
parent d4e3968e91b5860c698b700ed08a10e249bedbd6

Switch consteval to use llvm-cpu by default. (#14530)

This is possible because #14516 makes it possible to recursively target
LLVM-CPU.

Also added a check on the env var `IREE_COMPILER_DEBUG_CONSTEVAL` to
enable the debugging mode for consteval because these kinds of
compiler-in-compiler things are murderous to debug if needing to trace
through the whole thing and set a flag. I'll probably generalize this to
a more general "IREE_COMPILER_DEBUG" flag that we can use compiler-wide
for in the field debugging output (later).

compiler/src/iree/compiler/ConstEval/JitGlobals.cpp

diff --git a/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp b/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
index b682d2d..39f26ae 100644
--- a/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
+++ b/compiler/src/iree/compiler/ConstEval/JitGlobals.cpp
@@ -21,6 +21,8 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/SymbolTable.h"
 
+#include <cstdlib>
+
 #define DEBUG_TYPE "iree-const-eval"
 using llvm::dbgs;
 
@@ -31,7 +33,7 @@
 static llvm::cl::opt<std::string> clJitTargetBackend(
     "iree-consteval-jit-target-backend",
     llvm::cl::desc("Overrides the target backend used for JIT'ing."),
-    llvm::cl::init("vmvx"));
+    llvm::cl::init(""));
 
 static llvm::cl::opt<bool> clEnableDebug(
     "iree-consteval-jit-debug",
@@ -43,6 +45,14 @@
 
 namespace {
 
+static bool isDebugEnabled() {
+  if (clEnableDebug)
+    return true;
+  if (std::getenv("IREE_COMPILER_DEBUG_CONSTEVAL"))
+    return true;
+  return false;
+}
+
 // These options structs are not copy-constructable so we have to allocate them
 // shared.
 // TODO: See if we can make them copyable?
@@ -310,7 +320,7 @@
       : options(std::make_shared<CompileOptions>()),
         compilePipeline("builtin.module") {
     // Detect backend.
-    requestedTargetBackend = clJitTargetBackend;
+    requestedTargetBackend = resolveTargetBackend(targetRegistry);
     hasRequestedTargetBackend =
         targetRegistry.getTargetBackend(requestedTargetBackend) != nullptr;
     options->executableOptions.targets.push_back(requestedTargetBackend);
@@ -340,6 +350,21 @@
     compilePipeline.getDependentDialects(registry);
   }
 
+  static std::string
+  resolveTargetBackend(const IREE::HAL::TargetBackendRegistry &targetRegistry) {
+    if (clJitTargetBackend.empty()) {
+      // Default - choose something we have.
+      // First llvm-cpu then vmvx.
+      if (targetRegistry.getTargetBackend("llvm-cpu")) {
+        return std::string("llvm-cpu");
+      } else {
+        return std::string("vmvx");
+      }
+    }
+
+    return clJitTargetBackend;
+  }
+
   const SupportedFeatures getSupportedFeatures(MLIRContext *context) {
     SupportedFeatures s;
     Builder b(context);
@@ -374,7 +399,7 @@
     // Process each function through the runtime.
     for (JitFunctionDesc &jitFunction : jitFunctions) {
       std::optional<llvm::Timer> invokeTimer;
-      if (clEnableDebug) {
+      if (debugEnabled) {
         std::string timerName("Invoke ");
         timerName.append(jitFunction.name);
         invokeTimer.emplace(timerName, timerName, tg);
@@ -427,7 +452,7 @@
         }
       }
 
-      if (clEnableDebug) {
+      if (debugEnabled) {
         invokeTimer->stopTimer();
       }
     }
@@ -485,6 +510,8 @@
 
       if (succeeded(programBuilder.importInitializer(initOp))) {
         deadInitOps.push_back(initOp);
+      } else if (debugEnabled) {
+        dbgs() << "::: Rejected consteval initializer:\n" << initOp << "\n";
       }
     }
     if (programBuilder.getJitFunctions().empty()) {
@@ -493,9 +520,9 @@
     }
 
     std::optional<llvm::Timer> compileTimer;
-    if (clEnableDebug) {
-      dbgs() << "::: COMPILING JIT: " << programBuilder.getTargetModule()
-             << "\n";
+    if (debugEnabled) {
+      dbgs() << "::: COMPILING JIT (" << requestedTargetBackend
+             << "): " << programBuilder.getTargetModule() << "\n";
       compileTimer.emplace("iree-consteval-jit-compile", "Compiling", tg);
       compileTimer->startTimer();
     }
@@ -508,7 +535,7 @@
     if (failed(binary.translateFromModule(programBuilder.getTargetModule()))) {
       return signalPassFailure();
     }
-    if (clEnableDebug) {
+    if (debugEnabled) {
       compileTimer->stopTimer();
     }
 
@@ -536,6 +563,7 @@
   std::string requestedTargetBackend;
   std::shared_ptr<IREE::HAL::TargetBackend> targetBackend;
   bool hasRequestedTargetBackend;
+  bool debugEnabled = isDebugEnabled();
 };
 
 } // namespace