Revert pipelines (#948)

* Revert "[SD] Modify the flags to use --iree-preprocessing-pass-pipeline (#914)"

This reverts commit a783c089a9.

* Revert "Fix iree flags due to the change in shark-runtime (#944)"

This reverts commit 1d38d49162.

apps/stable_diffusion/src/models/opt_params.py

@@ -1,11 +1,6 @@
 import sys
 from transformers import CLIPTokenizer
-from apps.stable_diffusion.src.utils import (
-    models_db,
-    args,
-    get_shark_model,
-    get_opt_flags,
-)
+from apps.stable_diffusion.src.utils import models_db, args, get_shark_model
 
 
 hf_model_variant_map = {
@@ -24,14 +19,47 @@ def get_variant_version(hf_model_id):
 
 
 def get_params(bucket_key, model_key, model, is_tuned, precision):
+    iree_flags = []
+    if len(args.iree_vulkan_target_triple) > 0:
+        iree_flags.append(
+            f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
+        )
+
+    # Disable bindings fusion to work with moltenVK.
+    if sys.platform == "darwin":
+        iree_flags.append("-iree-stream-fuse-binding=false")
+
     try:
         bucket = models_db[0][bucket_key]
         model_name = models_db[1][model_key]
+        iree_flags += models_db[2][model][is_tuned][precision][
+            "default_compilation_flags"
+        ]
     except KeyError:
         raise Exception(
             f"{bucket_key}/{model_key} is not present in the models database"
         )
-    iree_flags = get_opt_flags(model, precision="fp16")
+
+    if (
+        "specified_compilation_flags"
+        in models_db[2][model][is_tuned][precision]
+    ):
+        device = (
+            args.device
+            if "://" not in args.device
+            else args.device.split("://")[0]
+        )
+        if (
+            device
+            not in models_db[2][model][is_tuned][precision][
+                "specified_compilation_flags"
+            ]
+        ):
+            device = "default_device"
+        iree_flags += models_db[2][model][is_tuned][precision][
+            "specified_compilation_flags"
+        ][device]
+
     return bucket, model_name, iree_flags
 
 

apps/stable_diffusion/src/utils/resources/model_db.json

@@ -78,5 +78,100 @@
     "dreamlike/v2_1base/vae/fp32/length_77/untuned":"dl_vae_23dec_fp32",
     "dreamlike/v2_1base/vae/fp32/length_77/untuned/base":"dl_vaebase_23dec_fp32",
     "dreamlike/v2_1base/clip/fp32/length_77/untuned":"dl_clip_23dec_fp32_77"
+  },
+  {
+    "unet": {
+      "tuned": {
+        "fp16": {
+          "default_compilation_flags": []
+        },
+        "fp32": {
+          "default_compilation_flags": []
+        }
+      },
+      "untuned": {
+        "fp16": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32"
+          ],
+          "specified_compilation_flags": {
+            "cuda": ["--iree-flow-enable-conv-nchw-to-nhwc-transform"],
+            "default_device": ["--iree-flow-enable-conv-img2col-transform"]
+          }
+        },
+        "fp32": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=16"
+          ]
+        }
+      }
+    },
+    "vae": {
+      "tuned": {
+        "fp16": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32",
+            "--iree-flow-enable-conv-img2col-transform"
+          ]
+        },
+        "fp32": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32",
+            "--iree-flow-enable-conv-img2col-transform"
+          ]
+        }
+      },
+      "untuned": {
+        "fp16": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32",
+            "--iree-flow-enable-conv-img2col-transform"
+          ]
+        },
+        "fp32": {
+          "default_compilation_flags": [
+            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=16"
+          ]
+        }
+      }
+    },
+    "clip": {
+      "tuned": {
+        "fp16": {
+          "default_compilation_flags": [
+            "--iree-flow-linalg-ops-padding-size=16",
+            "--iree-flow-enable-padding-linalg-ops"
+          ]
+        },
+        "fp32": {
+          "default_compilation_flags": [
+            "--iree-flow-linalg-ops-padding-size=16",
+            "--iree-flow-enable-padding-linalg-ops"
+          ]
+        }
+      },
+      "untuned": {
+        "fp16": {
+          "default_compilation_flags": [
+            "--iree-flow-linalg-ops-padding-size=16",
+            "--iree-flow-enable-padding-linalg-ops"
+          ]
+        },
+        "fp32": {
+          "default_compilation_flags": [
+            "--iree-flow-linalg-ops-padding-size=16",
+            "--iree-flow-enable-padding-linalg-ops"
+          ]
+        }
+      }
+    }
   }
 ]

apps/stable_diffusion/src/utils/resources/opt_flags.json

@@ -11,12 +11,19 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
-        ]
+          "--iree-flow-enable-padding-linalg-ops",
+          "--iree-flow-linalg-ops-padding-size=32"
+        ],
+        "specified_compilation_flags": {
+          "cuda": ["--iree-flow-enable-conv-nchw-to-nhwc-transform"],
+          "default_device": ["--iree-flow-enable-conv-img2col-transform"]
+        }
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+          "--iree-flow-enable-padding-linalg-ops",
+          "--iree-flow-linalg-ops-padding-size=16"
         ]
       }
     }
@@ -27,9 +34,9 @@
         "default_compilation_flags": [],
         "specified_compilation_flags": {
           "cuda": [],
-          "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
-          ]
+          "default_device": ["--iree-flow-enable-padding-linalg-ops",
+                             "--iree-flow-linalg-ops-padding-size=32",
+                             "--iree-flow-enable-conv-img2col-transform"]
         }
       },
       "fp32": {
@@ -37,7 +44,9 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32",
+            "--iree-flow-enable-conv-img2col-transform"
           ]
         }
       }
@@ -45,12 +54,16 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
+          "--iree-flow-enable-padding-linalg-ops",
+          "--iree-flow-linalg-ops-padding-size=32",
+          "--iree-flow-enable-conv-img2col-transform"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+          "--iree-flow-enable-padding-linalg-ops",
+          "--iree-flow-linalg-ops-padding-size=16"
         ]
       }
     }
@@ -59,24 +72,28 @@
     "tuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-linalg-ops-padding-size=16",
+          "--iree-flow-enable-padding-linalg-ops"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-linalg-ops-padding-size=16",
+          "--iree-flow-enable-padding-linalg-ops"
         ]
       }
     },
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-linalg-ops-padding-size=16",
+          "--iree-flow-enable-padding-linalg-ops"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-flow-linalg-ops-padding-size=16",
+          "--iree-flow-enable-padding-linalg-ops"
         ]
       }
     }

apps/stable_diffusion/src/utils/sd_annotation.py

@@ -26,7 +26,7 @@ def load_model_from_tank():
         get_variant_version,
     )
 
-    variant, version = get_variant_version(args.hf_model_id)
+    version, variant = get_variant_version(args.hf_model_id)
 
     shark_args.local_tank_cache = args.local_tank_cache
     bucket_key = f"{variant}/untuned"
@@ -62,7 +62,7 @@ def load_winograd_configs():
 def load_lower_configs():
     from apps.stable_diffusion.src.models import get_variant_version
 
-    variant, version = get_variant_version(args.hf_model_id)
+    version, variant = get_variant_version(args.hf_model_id)
 
     config_bucket = "gs://shark_tank/sd_tuned/configs/"
     config_version = version
@@ -109,19 +109,8 @@ def annotate_with_lower_configs(
 ):
     if use_winograd:
         dump_after = "iree-linalg-ext-convert-conv2d-to-winograd"
-        preprocess_flag = (
-            "--iree-preprocessing-pass-pipeline='builtin.module"
-            "(func.func(iree-preprocessing-convert-conv2d-to-img2col,"
-            "iree-preprocessing-pad-linalg-ops{pad-size=32},"
-            "iree-linalg-ext-convert-conv2d-to-winograd))' "
-        )
     else:
-        dump_after = "iree-preprocessing-pad-linalg-ops"
-        preprocess_flag = (
-            "--iree-preprocessing-pass-pipeline='builtin.module"
-            "(func.func(iree-preprocessing-convert-conv2d-to-img2col,"
-            "iree-preprocessing-pad-linalg-ops{pad-size=32}))' "
-        )
+        dump_after = "iree-flow-pad-linalg-ops"
 
     # Dump IR after padding/img2col/winograd passes
     device_spec_args = ""
@@ -143,9 +132,11 @@ def annotate_with_lower_configs(
         "--iree-input-type=tm_tensor "
         f"--iree-hal-target-backends={iree_target_map(device)} "
         f"{device_spec_args}"
-        f"{preprocess_flag}"
         "--iree-stream-resource-index-bits=64 "
         "--iree-vm-target-index-bits=64 "
+        "--iree-flow-enable-padding-linalg-ops "
+        "--iree-flow-linalg-ops-padding-size=32 "
+        "--iree-flow-enable-conv-img2col-transform "
         f"--mlir-print-ir-after={dump_after} "
         "--compile-to=flow "
         f"2>{args.annotation_output}/dump_after_winograd.mlir "

shark/examples/shark_inference/upscaler/opt_params.py

@@ -13,15 +13,20 @@ if BATCH_SIZE != 1:
 
 
 unet_flag = [
-    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
+    "--iree-flow-enable-padding-linalg-ops",
+    "--iree-flow-linalg-ops-padding-size=32",
+    "--iree-flow-enable-conv-img2col-transform",
 ]
 
 vae_flag = [
-    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-convert-conv-nchw-to-nhwc,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+    "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+    "--iree-flow-enable-padding-linalg-ops",
+    "--iree-flow-linalg-ops-padding-size=16",
 ]
 
 clip_flag = [
-    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+    "--iree-flow-linalg-ops-padding-size=16",
+    "--iree-flow-enable-padding-linalg-ops",
 ]
 
 bucket = "gs://shark_tank/stable_diffusion/"

shark/iree_utils/compile_utils.py

@@ -80,17 +80,11 @@ def get_iree_common_args():
 def get_model_specific_args():
     ms_args = []
     if shark_args.enable_conv_transform == True:
-        ms_args += [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-convert-conv-nchw-to-nhwc))"
-        ]
+        ms_args += ["--iree-flow-enable-conv-nchw-to-nhwc-transform"]
     if shark_args.enable_img2col_transform == True:
-        ms_args += [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col))"
-        ]
+        ms_args += ["--iree-flow-enable-conv-img2col-transform"]
     if shark_args.use_winograd == True:
-        ms_args += [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-linalg-ext-convert-conv2d-to-winograd))"
-        ]
+        ms_args += ["--iree-flow-enable-conv-winograd-transform"]
     return ms_args