Add flag for toggling constant folding. (#1714)

nod-ai · Aug 4, 2023 · c9de272 · c9de272
1 parent a5b13fc
commit c9de272
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py
@@ -519,6 +519,12 @@ def is_valid_file(arg):
     "in shark importer. Does nothing if import_mlir is false (the default).",
 )
 
+p.add_argument(
+    "--iree_constant_folding",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="Controls constant folding in iree-compile for all SD models.",
+)
 
 ##############################################################################
 # Web UI flags

diff --git a/apps/stable_diffusion/src/utils/utils.py b/apps/stable_diffusion/src/utils/utils.py
@@ -500,6 +500,12 @@ def get_opt_flags(model, precision="fp16"):
             f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
         )
 
+    if args.iree_constant_folding == False:
+        iree_flags.append("--iree-opt-const-expr-hoisting=False")
+        iree_flags.append(
+            "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807"
+        )
+
     # Disable bindings fusion to work with moltenVK.
     if sys.platform == "darwin":
         iree_flags.append("-iree-stream-fuse-binding=false")