Rework consteval to be more memory efficient. (#14504)

* Also adds an option to use the llvm-cpu backend instead of VMVX. This
isn't quite ready for primetime because I need to isolate some flags
better, but it can be used in a pinch.
* Changes the way that JIT programs are constructed so that stateless
functions are compiled at once for all constants and the data is passed
in/out vs being cloned from module globals.
* Each global has a public function to compute it, allowing us to run
them one by one or in batches to keep memory use under control.
* In this patch, I am just running them one by one and also haven't
optimized memory by use of resources yet. Still a big improvement.
* Drops compilation time of h2ogpt from ~45m to 2m47s. More optimization
is possible.
* JIT'ing globals now takes 52s for this model (vs ~42m with the Linalg
const evaler).
* Memory pressure is kept under control and does not regress from the
current state (more optimization is possible - just a starting point).

I also changed the behavior of the constexpr hoisting and jit passes to
produce/process initializers with the `iree.compiler.consteval`
attribute. This was advisable because in the new way of doing
evaluation, we are doing a non-general transformation on the
initializers and I didn't want it stumbling over arbitrary initializers
in the wild. Having the behavior be opt-in seemed prudent.

Flag changes:

* `iree-consteval-jit-use-vmvx` (default true): Uses the VMVX backend.
When false, uses the LLVMCPU backend. I'll be doing work to change this
to false by default when LLVMCPU is available.
* `iree-consteval-jit-debug`: Prints debugging information about
constant evaluation.
* `iree-opt-const-eval`: Flipped to true since it now only processes
initializers targeted at it and is therefore safe to always have
enabled.

compiler/src/iree/compiler/API/Internal/Embed.cpp

@@ -581,9 +581,11 @@ Invocation::Invocation(Session &session)
   // Since the jitter invokes much of the top-level compiler recursively,
   // it must be injected at the top-level here vs in the pass pipeline
   // (or else the circular dependency cannot be resolved).
-  pipelineHooks.buildConstEvalPassPipelineCallback = [](OpPassManager &pm) {
-    pm.addPass(ConstEval::createJitGlobalsPass());
-  };
+  auto &targetRegistry = session.targetRegistry;
+  pipelineHooks.buildConstEvalPassPipelineCallback =
+      [&targetRegistry](OpPassManager &pm) {
+        pm.addPass(ConstEval::createJitGlobalsPass(targetRegistry));
+      };
   // The PluginSession implements PipelineExtensions and delegates it to
   // activated plugins.
   pipelineHooks.pipelineExtensions = &session.pluginSession;

compiler/src/iree/compiler/ConstEval/BUILD.bazel

@@ -57,6 +57,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Pipelines",
         "//compiler/src/iree/compiler/Utils",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
@@ -79,6 +80,7 @@ iree_compiler_cc_library(
         "//runtime/src/iree/tooling:vm_util",
         "//runtime/src/iree/vm",
         "//runtime/src/iree/vm/bytecode:module",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
     ],
 )

compiler/src/iree/compiler/ConstEval/CMakeLists.txt

@@ -46,6 +46,7 @@ iree_cc_library(
     ::PassesIncGen
     ::Runtime
     LLVMSupport
+    MLIRArithDialect
     MLIRFuncDialect
     MLIRIR
     MLIRPass
@@ -62,6 +63,7 @@ iree_cc_library(
   SRCS
     "Runtime.cpp"
   DEPS
+    LLVMSupport
     MLIRIR
     iree::compiler::Dialect::VM::Target::Bytecode
     iree::hal