Add init_at_opt

Signed-off-by: Ian Wood <[email protected]>
iree-org · Feb 27, 2025 · 46d1505 · 46d1505
1 parent 396cd99
commit 46d1505
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 40 deletions.
diff --git a/compiler/src/iree/compiler/Pipelines/Options.cpp b/compiler/src/iree/compiler/Pipelines/Options.cpp
@@ -136,17 +136,14 @@ void GlobalOptimizationOptions::applyOptimization(
     const OptionsBinder &binder, const GlobalPipelineOptions &globalLevel) {
   binder.overrideDefault("iree-global-optimization-opt-level", optLevel,
                          globalLevel.optLevel);
-
-  if (optLevel != llvm::OptimizationLevel::O0) {
-    binder.overrideDefault("iree-opt-strip-assertions", stripAssertions, true);
-    llvm::dbgs() << "stripAssertions " << stripAssertions << '\n';
-  }
+  binder.applyOptimization("iree-opt-strip-assertions", stripAssertions,
+                           optLevel);
 };
 
 void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
   static llvm::cl::OptionCategory category(
       "IREE options for controlling global optimizations.");
-  binder.opt<llvm::OptimizationLevel>(
+  binder.optimizationLevel(
       "iree-global-optimization-opt-level", optLevel,
       llvm::cl::desc("Optimization level for the this pipeline"),
       llvm::cl::cat(category));
@@ -187,6 +184,7 @@ void GlobalOptimizationOptions::bindOptions(OptionsBinder &binder) {
           "Reduces numeric precision to lower bit depths where possible."),
       llvm::cl::cat(category));
   binder.opt<bool>("iree-opt-strip-assertions", stripAssertions,
+                   init_at_opt(llvm::OptimizationLevel::O1, true),
                    llvm::cl::desc("Strips debug assertions after any useful "
                                   "information has been extracted."),
                    llvm::cl::cat(category));

diff --git a/compiler/src/iree/compiler/Utils/OptionUtils.h b/compiler/src/iree/compiler/Utils/OptionUtils.h
@@ -18,22 +18,23 @@
 
 namespace mlir::iree_compiler {
 
+struct opt_initializer_base {};
+
 template <typename Ty>
-struct opt_initializer {
-  const Ty &init;
+struct opt_initializer : opt_initializer_base {
+  const Ty init;
   llvm::OptimizationLevel optLevel;
   opt_initializer(const llvm::OptimizationLevel opt, const Ty &val)
       : init(val), optLevel(opt) {}
-
-  void apply(llvm::OptimizationLevel opt, Ty &val) const {
-    if (opt == optLevel) {
+  void apply(const llvm::OptimizationLevel inLevel, Ty &val) const {
+    assert(inLevel.getSizeLevel() == 0 && "size level not implemented");
+    if (inLevel.getSpeedupLevel() >= optLevel.getSpeedupLevel())
       val = init;
-    }
   }
 };
 
-struct OptScope {};
-
+/// Initialize the value of a variable if the optimization level is at least
+/// the specified level.
 template <typename Ty>
 opt_initializer<Ty> init_at_opt(llvm::OptimizationLevel optLevel,
                                 const Ty &val) {
@@ -98,35 +99,35 @@ class OptionsBinder {
             typename... Mods>
   void opt(llvm::StringRef name, V &value, opt_initializer<K> init,
            opt_initializer<Ks>... inits, Mods... Ms) {
-    auto [changedCallback, clCallback] = makeChangedCallback<V>();
-    if (!scope) {
-      // Bind global options.
-      auto opt = std::make_unique<llvm::cl::opt<T, /*ExternalStorage=*/true>>(
-          name, llvm::cl::location(value), llvm::cl::init(value), clCallback,
-          std::forward<Mods>(Ms)...);
-      auto defaultCallback = makeDefaultCallback(&value);
-      getOptionsStorage()[name] = OptionInfo{std::move(opt), /*print=*/nullptr,
-                                             /*isChanged=*/changedCallback,
-                                             /*isDefault*/ defaultCallback};
-    } else {
-      // Bind local options.
-      auto option =
-          std::make_unique<llvm::cl::opt<T, /*ExternalStorage=*/true>>(
-              name, llvm::cl::sub(*scope), llvm::cl::location(value),
-              llvm::cl::init(value), clCallback, std::forward<Mods>(Ms)...);
-      auto printCallback =
-          makePrintCallback(option->ArgStr, option->getParser(), &value);
-      auto defaultCallback = makeDefaultCallback(&value);
-      getOptionsStorage()[name] = OptionInfo{
-          std::move(option), /*print=*/printCallback,
-          /*isChanged=*/changedCallback, /*isDefault*/ defaultCallback};
-    }
+    opt<V>(name, value, Ms...);
+    getOptionsStorage()[name].setOptLevels.push_back(
+        std::make_unique<opt_initializer<K>>(init));
+    (getOptionsStorage()[name].setOptLevels.push_back(
+         std::make_unique<opt_initializer<Ks>>(inits)),
+     ...);
   }
 
   template <typename... Mods>
-  OptScope optimizationLevel(llvm::StringRef name,
-                             llvm::OptimizationLevel &value, Mods... mods) {
-    return {};
+  void optimizationLevel(llvm::StringRef name, llvm::OptimizationLevel &value,
+                         Mods... Ms) {
+    opt<llvm::OptimizationLevel>(name, value, Ms...);
+  }
+
+  template <typename T>
+  void applyOptimization(llvm::StringRef name, T &value,
+                         llvm::OptimizationLevel optLevel) const {
+    const auto infoIt = getOptionsStorage().find(name);
+    assert(infoIt != getOptionsStorage().end() && "Option not found");
+    auto changedCallback = infoIt->getSecond().isChanged;
+    assert(changedCallback && "Expected changed callback");
+    if (changedCallback()) {
+      return;
+    }
+    const auto &setOptLevels = infoIt->getSecond().setOptLevels;
+    for (const auto &init : setOptLevels) {
+      reinterpret_cast<opt_initializer<T> *>(init.get())
+          ->apply(optLevel, value);
+    }
   }
 
   bool isFlagSet(llvm::StringRef name) const {
@@ -196,6 +197,8 @@ class OptionsBinder {
     PrintCallback print;
     ChangedCallback isChanged;
     DefaultCallback isDefault;
+
+    llvm::SmallVector<std::unique_ptr<opt_initializer_base>> setOptLevels;
   };
   using OptionsStorage = llvm::DenseMap<llvm::StringRef, OptionInfo>;