From c9de2729b22a4ffce4a484ed6b38dce30886cb51 Mon Sep 17 00:00:00 2001
From: Ean Garvey <87458719+monorimet@users.noreply.github.com>
Date: Fri, 4 Aug 2023 06:55:52 -0500
Subject: [PATCH] Add flag for toggling constant folding. (#1714)

---
 apps/stable_diffusion/src/utils/stable_args.py | 6 ++++++
 apps/stable_diffusion/src/utils/utils.py       | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py
index c2601fbbd4..202d9047da 100644
--- 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
index 58a0ff5b83..c7814c6ed5 100644
--- 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")