Revert "Revert pipelines (#948)" (#951)

This reverts commit 8115b26079.
Additionally fixes img2col by adding detach elementwise from named op
passes.

apps/stable_diffusion/src/models/opt_params.py

@@ -1,6 +1,11 @@
 import sys
 from transformers import CLIPTokenizer
-from apps.stable_diffusion.src.utils import models_db, args, get_shark_model
+from apps.stable_diffusion.src.utils import (
+    models_db,
+    args,
+    get_shark_model,
+    get_opt_flags,
+)
 
 
 hf_model_variant_map = {
@@ -19,47 +24,14 @@ 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"
         )
-
-    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]
-
+    iree_flags = get_opt_flags(model, precision="fp16")
     return bucket, model_name, iree_flags
 
 

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

@@ -78,100 +78,5 @@
     "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,19 +11,12 @@
     "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"]
-        }
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
+        ]
       },
       "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"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }
@@ -34,9 +27,9 @@
         "default_compilation_flags": [],
         "specified_compilation_flags": {
           "cuda": [],
-          "default_device": ["--iree-flow-enable-padding-linalg-ops",
-                             "--iree-flow-linalg-ops-padding-size=32",
-                             "--iree-flow-enable-conv-img2col-transform"]
+          "default_device": [
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
+          ]
         }
       },
       "fp32": {
@@ -44,9 +37,7 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-flow-enable-padding-linalg-ops",
-            "--iree-flow-linalg-ops-padding-size=32",
-            "--iree-flow-enable-conv-img2col-transform"
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
           ]
         }
       }
@@ -54,16 +45,12 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-flow-enable-padding-linalg-ops",
-          "--iree-flow-linalg-ops-padding-size=32",
-          "--iree-flow-enable-conv-img2col-transform"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
         ]
       },
       "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"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }
@@ -72,28 +59,24 @@
     "tuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-flow-linalg-ops-padding-size=16",
-          "--iree-flow-enable-padding-linalg-ops"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-flow-linalg-ops-padding-size=16",
-          "--iree-flow-enable-padding-linalg-ops"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     },
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-flow-linalg-ops-padding-size=16",
-          "--iree-flow-enable-padding-linalg-ops"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-flow-linalg-ops-padding-size=16",
-          "--iree-flow-enable-padding-linalg-ops"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }

apps/stable_diffusion/src/utils/sd_annotation.py

@@ -26,7 +26,7 @@ def load_model_from_tank():
         get_variant_version,
     )
 
-    version, variant = get_variant_version(args.hf_model_id)
+    variant, version = 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
 
-    version, variant = get_variant_version(args.hf_model_id)
+    variant, version = get_variant_version(args.hf_model_id)
 
     config_bucket = "gs://shark_tank/sd_tuned/configs/"
     config_version = version
@@ -109,8 +109,19 @@ 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-flow-pad-linalg-ops"
+        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 IR after padding/img2col/winograd passes
     device_spec_args = ""
@@ -132,11 +143,9 @@ 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,20 +13,15 @@ if BATCH_SIZE != 1:
 
 
 unet_flag = [
-    "--iree-flow-enable-padding-linalg-ops",
-    "--iree-flow-linalg-ops-padding-size=32",
-    "--iree-flow-enable-conv-img2col-transform",
+    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
 ]
 
 vae_flag = [
-    "--iree-flow-enable-conv-nchw-to-nhwc-transform",
-    "--iree-flow-enable-padding-linalg-ops",
-    "--iree-flow-linalg-ops-padding-size=16",
+    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-convert-conv-nchw-to-nhwc,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
 ]
 
 clip_flag = [
-    "--iree-flow-linalg-ops-padding-size=16",
-    "--iree-flow-enable-padding-linalg-ops",
+    "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
 ]
 
 bucket = "gs://shark_tank/stable_diffusion/"

shark/iree_utils/compile_utils.py

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