Skip to content

Commit a783c08

Browse files
yzhang93yzhang93
authored
[SD] Modify the flags to use --iree-preprocessing-pass-pipeline (huggingface#914)
* [SD] Modify the flags to use --iree-preprocessing-pass-pipeline * Fix flags in sd_annotation
1 parent e7907dc commit a783c08

File tree

4 files changed

+35
-166
lines changed

4 files changed

+35
-166
lines changed

‎apps/stable_diffusion/src/models/opt_params.py‎

Lines changed: 7 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,11 @@
11
import sys
22
from transformers import CLIPTokenizer
3-
from apps.stable_diffusion.src.utils import models_db, args, get_shark_model
3+
from apps.stable_diffusion.src.utils import (
4+
models_db,
5+
args,
6+
get_shark_model,
7+
get_opt_flags,
8+
)
49

510

611
hf_model_variant_map = {
@@ -19,47 +24,14 @@ def get_variant_version(hf_model_id):
1924

2025

2126
def get_params(bucket_key, model_key, model, is_tuned, precision):
22-
iree_flags = []
23-
if len(args.iree_vulkan_target_triple) > 0:
24-
iree_flags.append(
25-
f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
26-
)
27-
28-
# Disable bindings fusion to work with moltenVK.
29-
if sys.platform == "darwin":
30-
iree_flags.append("-iree-stream-fuse-binding=false")
31-
3227
try:
3328
bucket = models_db[0][bucket_key]
3429
model_name = models_db[1][model_key]
35-
iree_flags += models_db[2][model][is_tuned][precision][
36-
"default_compilation_flags"
37-
]
3830
except KeyError:
3931
raise Exception(
4032
f"{bucket_key}/{model_key} is not present in the models database"
4133
)
42-
43-
if (
44-
"specified_compilation_flags"
45-
in models_db[2][model][is_tuned][precision]
46-
):
47-
device = (
48-
args.device
49-
if "://" not in args.device
50-
else args.device.split("://")[0]
51-
)
52-
if (
53-
device
54-
not in models_db[2][model][is_tuned][precision][
55-
"specified_compilation_flags"
56-
]
57-
):
58-
device = "default_device"
59-
iree_flags += models_db[2][model][is_tuned][precision][
60-
"specified_compilation_flags"
61-
][device]
62-
34+
iree_flags = get_opt_flags(model, precision="fp16")
6335
return bucket, model_name, iree_flags
6436

6537

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

Lines changed: 0 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -78,100 +78,5 @@
7878
"dreamlike/v2_1base/vae/fp32/length_77/untuned":"dl_vae_23dec_fp32",
7979
"dreamlike/v2_1base/vae/fp32/length_77/untuned/base":"dl_vaebase_23dec_fp32",
8080
"dreamlike/v2_1base/clip/fp32/length_77/untuned":"dl_clip_23dec_fp32_77"
81-
},
82-
{
83-
"unet": {
84-
"tuned": {
85-
"fp16": {
86-
"default_compilation_flags": []
87-
},
88-
"fp32": {
89-
"default_compilation_flags": []
90-
}
91-
},
92-
"untuned": {
93-
"fp16": {
94-
"default_compilation_flags": [
95-
"--iree-flow-enable-padding-linalg-ops",
96-
"--iree-flow-linalg-ops-padding-size=32"
97-
],
98-
"specified_compilation_flags": {
99-
"cuda": ["--iree-flow-enable-conv-nchw-to-nhwc-transform"],
100-
"default_device": ["--iree-flow-enable-conv-img2col-transform"]
101-
}
102-
},
103-
"fp32": {
104-
"default_compilation_flags": [
105-
"--iree-flow-enable-conv-nchw-to-nhwc-transform",
106-
"--iree-flow-enable-padding-linalg-ops",
107-
"--iree-flow-linalg-ops-padding-size=16"
108-
]
109-
}
110-
}
111-
},
112-
"vae": {
113-
"tuned": {
114-
"fp16": {
115-
"default_compilation_flags": [
116-
"--iree-flow-enable-padding-linalg-ops",
117-
"--iree-flow-linalg-ops-padding-size=32",
118-
"--iree-flow-enable-conv-img2col-transform"
119-
]
120-
},
121-
"fp32": {
122-
"default_compilation_flags": [
123-
"--iree-flow-enable-padding-linalg-ops",
124-
"--iree-flow-linalg-ops-padding-size=32",
125-
"--iree-flow-enable-conv-img2col-transform"
126-
]
127-
}
128-
},
129-
"untuned": {
130-
"fp16": {
131-
"default_compilation_flags": [
132-
"--iree-flow-enable-padding-linalg-ops",
133-
"--iree-flow-linalg-ops-padding-size=32",
134-
"--iree-flow-enable-conv-img2col-transform"
135-
]
136-
},
137-
"fp32": {
138-
"default_compilation_flags": [
139-
"--iree-flow-enable-conv-nchw-to-nhwc-transform",
140-
"--iree-flow-enable-padding-linalg-ops",
141-
"--iree-flow-linalg-ops-padding-size=16"
142-
]
143-
}
144-
}
145-
},
146-
"clip": {
147-
"tuned": {
148-
"fp16": {
149-
"default_compilation_flags": [
150-
"--iree-flow-linalg-ops-padding-size=16",
151-
"--iree-flow-enable-padding-linalg-ops"
152-
]
153-
},
154-
"fp32": {
155-
"default_compilation_flags": [
156-
"--iree-flow-linalg-ops-padding-size=16",
157-
"--iree-flow-enable-padding-linalg-ops"
158-
]
159-
}
160-
},
161-
"untuned": {
162-
"fp16": {
163-
"default_compilation_flags": [
164-
"--iree-flow-linalg-ops-padding-size=16",
165-
"--iree-flow-enable-padding-linalg-ops"
166-
]
167-
},
168-
"fp32": {
169-
"default_compilation_flags": [
170-
"--iree-flow-linalg-ops-padding-size=16",
171-
"--iree-flow-enable-padding-linalg-ops"
172-
]
173-
}
174-
}
175-
}
17681
}
17782
]

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

Lines changed: 13 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -11,19 +11,12 @@
1111
"untuned": {
1212
"fp16": {
1313
"default_compilation_flags": [
14-
"--iree-flow-enable-padding-linalg-ops",
15-
"--iree-flow-linalg-ops-padding-size=32"
16-
],
17-
"specified_compilation_flags": {
18-
"cuda": ["--iree-flow-enable-conv-nchw-to-nhwc-transform"],
19-
"default_device": ["--iree-flow-enable-conv-img2col-transform"]
20-
}
14+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
15+
]
2116
},
2217
"fp32": {
2318
"default_compilation_flags": [
24-
"--iree-flow-enable-conv-nchw-to-nhwc-transform",
25-
"--iree-flow-enable-padding-linalg-ops",
26-
"--iree-flow-linalg-ops-padding-size=16"
19+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
2720
]
2821
}
2922
}
@@ -34,36 +27,30 @@
3427
"default_compilation_flags": [],
3528
"specified_compilation_flags": {
3629
"cuda": [],
37-
"default_device": ["--iree-flow-enable-padding-linalg-ops",
38-
"--iree-flow-linalg-ops-padding-size=32",
39-
"--iree-flow-enable-conv-img2col-transform"]
30+
"default_device": [
31+
"--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))"
32+
]
4033
}
4134
},
4235
"fp32": {
4336
"default_compilation_flags": [],
4437
"specified_compilation_flags": {
4538
"cuda": [],
4639
"default_device": [
47-
"--iree-flow-enable-padding-linalg-ops",
48-
"--iree-flow-linalg-ops-padding-size=32",
49-
"--iree-flow-enable-conv-img2col-transform"
40+
"--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))"
5041
]
5142
}
5243
}
5344
},
5445
"untuned": {
5546
"fp16": {
5647
"default_compilation_flags": [
57-
"--iree-flow-enable-padding-linalg-ops",
58-
"--iree-flow-linalg-ops-padding-size=32",
59-
"--iree-flow-enable-conv-img2col-transform"
48+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
6049
]
6150
},
6251
"fp32": {
6352
"default_compilation_flags": [
64-
"--iree-flow-enable-conv-nchw-to-nhwc-transform",
65-
"--iree-flow-enable-padding-linalg-ops",
66-
"--iree-flow-linalg-ops-padding-size=16"
53+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
6754
]
6855
}
6956
}
@@ -72,28 +59,24 @@
7259
"tuned": {
7360
"fp16": {
7461
"default_compilation_flags": [
75-
"--iree-flow-linalg-ops-padding-size=16",
76-
"--iree-flow-enable-padding-linalg-ops"
62+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
7763
]
7864
},
7965
"fp32": {
8066
"default_compilation_flags": [
81-
"--iree-flow-linalg-ops-padding-size=16",
82-
"--iree-flow-enable-padding-linalg-ops"
67+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
8368
]
8469
}
8570
},
8671
"untuned": {
8772
"fp16": {
8873
"default_compilation_flags": [
89-
"--iree-flow-linalg-ops-padding-size=16",
90-
"--iree-flow-enable-padding-linalg-ops"
74+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
9175
]
9276
},
9377
"fp32": {
9478
"default_compilation_flags": [
95-
"--iree-flow-linalg-ops-padding-size=16",
96-
"--iree-flow-enable-padding-linalg-ops"
79+
"--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))"
9780
]
9881
}
9982
}

‎apps/stable_diffusion/src/utils/sd_annotation.py‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def load_model_from_tank():
2626
get_variant_version,
2727
)
2828

29-
version, variant = get_variant_version(args.hf_model_id)
29+
variant, version = get_variant_version(args.hf_model_id)
3030

3131
shark_args.local_tank_cache = args.local_tank_cache
3232
bucket_key = f"{variant}/untuned"
@@ -62,7 +62,7 @@ def load_winograd_configs():
6262
def load_lower_configs():
6363
from apps.stable_diffusion.src.models import get_variant_version
6464

65-
version, variant = get_variant_version(args.hf_model_id)
65+
variant, version = get_variant_version(args.hf_model_id)
6666

6767
config_bucket = "gs://shark_tank/sd_tuned/configs/"
6868
config_version = version
@@ -109,8 +109,19 @@ def annotate_with_lower_configs(
109109
):
110110
if use_winograd:
111111
dump_after = "iree-linalg-ext-convert-conv2d-to-winograd"
112+
preprocess_flag = (
113+
"--iree-preprocessing-pass-pipeline='builtin.module"
114+
"(func.func(iree-preprocessing-convert-conv2d-to-img2col,"
115+
"iree-preprocessing-pad-linalg-ops{pad-size=32},"
116+
"iree-linalg-ext-convert-conv2d-to-winograd))' "
117+
)
112118
else:
113-
dump_after = "iree-flow-pad-linalg-ops"
119+
dump_after = "iree-preprocessing-pad-linalg-ops"
120+
preprocess_flag = (
121+
"--iree-preprocessing-pass-pipeline='builtin.module"
122+
"(func.func(iree-preprocessing-convert-conv2d-to-img2col,"
123+
"iree-preprocessing-pad-linalg-ops{pad-size=32}))' "
124+
)
114125

115126
# Dump IR after padding/img2col/winograd passes
116127
device_spec_args = ""
@@ -132,11 +143,9 @@ def annotate_with_lower_configs(
132143
"--iree-input-type=tm_tensor "
133144
f"--iree-hal-target-backends={iree_target_map(device)} "
134145
f"{device_spec_args}"
146+
f"{preprocess_flag}"
135147
"--iree-stream-resource-index-bits=64 "
136148
"--iree-vm-target-index-bits=64 "
137-
"--iree-flow-enable-padding-linalg-ops "
138-
"--iree-flow-linalg-ops-padding-size=32 "
139-
"--iree-flow-enable-conv-img2col-transform "
140149
f"--mlir-print-ir-after={dump_after} "
141150
"--compile-to=flow "
142151
f"2>{args.annotation_output}/dump_after_winograd.mlir "

0 commit comments

Comments
 (0)