[SD3 Model] VAE + MMDIT dispatches

[yzhang93]

VAE:

Dispatch Type	Shape	Compilation Time [ms]	Execution Time [ms]
matmul	256x65536x512	11968	1233
matmul	128x262144x256	13050	1647
matmul_transpose_b	4096x512x512	1870	151
batch_matmul	1x4096x4096x512	19323	1544
batch_matmul	1x4096x512x4096	26673	1732
conv_2d_nchw_fchw	1x512x64x64x16x3x3
conv_2d_nchw_fchw	1x512x64x64x512x3x3
conv_2d_nchw_fchw	1x512x128x128x512x3x3
conv_2d_nchw_fchw	1x512x256x256x512x3x3
conv_2d_nchw_fchw	1x256x256x256x512x3x3
conv_2d_nchw_fchw	1x256x256x256x256x3x3
conv_2d_nchw_fchw	1x256x512x512x256x3x3
conv_2d_nchw_fchw	1x128x512x512x256x3x3
conv_2d_nchw_fchw	1x128x512x512x128x3x3
conv_2d_nchw_fchw	1x3x512x512x128x3x3

MMDIT:

Dispatch Type	Shape	Compilation Time [ms]	Execution Time [ms]
matmul_transpose_b	1x1536x256
matmul_transpose_b	1x1536x1536
matmul_transpose_b	2x1536x2048
matmul_transpose_b	2x1536x1536
matmul_transpose_b	2x3072x1536
matmul_transpose_b	2x9216x1536
matmul_transpose_b	308x1536x4096
matmul_transpose_b	308x1536x1536
matmul_transpose_b	2048x1536x1536
matmul_transpose_b	2048x6144x1536
matmul_transpose_b	2048x1536x6144
matmul_transpose_b	2048x64x1536
conv_2d_nchw_fchw	2x1536x32x32x16x2x2

Text Encoder:

Dispatch Type	Shape	Compilation Time [ms]	Execution Time [ms]
matmul_transpose_b	77x768x768
matmul_transpose_b	77x3072x768
matmul_transpose_b	77x768x3072
matmul_transpose_b	77x1280x1280
matmul_transpose_b	77x20x64
matmul_transpose_b	77x5120x1280
matmul_transpose_b	77x1280x5120
matmul_transpose_b	77x4096x4096
matmul_transpose_b	77x64x64
matmul_transpose_b	77x10240x4096
matmul_transpose_b	77x4096x10240
matmul_transpose_b	1x1280x1280
batch_matmul	12x77x77x64
batch_matmul	12x77x64x77
batch_matmul	20x77x77x64
batch_matmul	20x77x64x77
batch_matmul	64x77x77x64
batch_matmul	64x77x64x77

Note: Time is collected with objectFifo pipeline, and data type is bf16 -> f32.

Command:
iree-compile --mlir-elide-elementsattrs-if-larger=2 --iree-hal-target-backends=llvm-cpu stable_diffusion_3_medium_diffusers_bs1_512x512_fp32_vae_decomp_attn.mlir --iree-dispatch-creation-enable-aggressive-fusion --iree-opt-data-tiling=false --iree-dispatch-creation-enable-fuse-padding-into-linalg-consumer-ops -o vae.vmfb --iree-hal-dump-executable-sources-to=vae_dispathces

[jtuyls]

@yzhang93 the shapes are MxNxK for matmul, right?

[yzhang93]

@yzhang93 the shapes are MxNxK for matmul, right?

Yes, without specific notation it would always be MxNxK.