Add SDXL EulerDiscreteScheduler compilation test (iree-org#19315)

This commit adds a compilation test for the scheduler used in sdxl. We
can iterate on this and add some tests that runs the different functions
as well (`run_initialize`, `run_scale`, `run_step`) and checks for
accuracy.

ci-exactly: build_packages, regression_test

---------

Signed-off-by: saienduri <[email protected]>
Signed-off-by: Giacomo Serafini <[email protected]>

Author: saienduri

experimental/regression_suite/shark-test-suite-models/conftest.py

@@ -9,13 +9,15 @@ class VmfbManager:
     sdxl_clip_cpu_vmfb = None
     sdxl_vae_cpu_vmfb = None
     sdxl_unet_fp16_cpu_vmfb = None
+    sdxl_unet_fp16_cpu_pipeline_vmfb = None
+    sdxl_scheduler_cpu_vmfb = None
     sdxl_clip_rocm_vmfb = None
     sdxl_vae_rocm_vmfb = None
     sdxl_unet_fp16_rocm_vmfb = None
     sdxl_punet_int8_fp16_rocm_vmfb = None
     sdxl_punet_int8_fp8_rocm_vmfb = None
-    sdxl_unet_fp16_cpu_pipeline_vmfb = None
     sdxl_unet_fp16_rocm_pipeline_vmfb = None
+    sdxl_scheduler_rocm_vmfb = None
     sd3_clip_cpu_vmfb = None
     sd3_vae_cpu_vmfb = None
     sd3_mmdit_cpu_vmfb = None

experimental/regression_suite/shark-test-suite-models/sd3/test_clip.py

@@ -10,7 +10,7 @@
 from conftest import VmfbManager
 from pathlib import Path
 
-rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
 vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
 
 ###############################################################################

experimental/regression_suite/shark-test-suite-models/sd3/test_mmdit.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 from conftest import VmfbManager
 
-rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
 vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
 
 ###############################################################################

experimental/regression_suite/shark-test-suite-models/sd3/test_vae.py

@@ -10,7 +10,7 @@
 from conftest import VmfbManager
 from pathlib import Path
 
-rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
 vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
 
 ###############################################################################

experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py

@@ -10,7 +10,7 @@
 from conftest import VmfbManager
 from pathlib import Path
 
-rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
 vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
 
 ###############################################################################

experimental/regression_suite/shark-test-suite-models/sdxl/test_scheduler.py

@@ -0,0 +1,76 @@
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import pytest
+from ireers_tools import *
+import os
+from conftest import VmfbManager
+from pathlib import Path
+
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
+vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
+
+###############################################################################
+# Fixtures
+###############################################################################
+
+sdxl_scheduler_mlir = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sdxl-scheduler/11-26-2024/model.mlir",
+    group="sdxl_scheduler",
+)
+
+CPU_COMPILE_FLAGS = [
+    "--iree-hal-target-backends=llvm-cpu",
+    "--iree-llvmcpu-target-cpu-features=host",
+    "--iree-llvmcpu-fail-on-out-of-bounds-stack-allocation=false",
+    "--iree-llvmcpu-distribution-size=32",
+    "--iree-opt-const-eval=false",
+    "--iree-opt-strip-assertions=true",
+    "--iree-llvmcpu-enable-ukernels=all",
+    "--iree-global-opt-enable-quantized-matmul-reassociation",
+]
+
+
+ROCM_COMPILE_FLAGS = [
+    "--iree-hal-target-backends=rocm",
+    f"--iree-hip-target={rocm_chip}",
+    "--iree-opt-const-eval=false",
+    "--iree-global-opt-propagate-transposes=true",
+    "--iree-llvmgpu-enable-prefetch=true",
+    "--iree-execution-model=async-external",
+    "--iree-preprocessing-pass-pipeline=builtin.module(iree-preprocessing-transpose-convolution-pipeline,iree-preprocessing-pad-to-intrinsics)",
+    "--iree-scheduling-dump-statistics-format=json",
+    "--iree-scheduling-dump-statistics-file=compilation_info.json",
+]
+
+###############################################################################
+# CPU
+###############################################################################
+
+
+def test_compile_scheduler_cpu(sdxl_scheduler_mlir):
+    VmfbManager.sdxl_scheduler_cpu_vmfb = iree_compile(
+        sdxl_scheduler_mlir,
+        CPU_COMPILE_FLAGS,
+        Path(vmfb_dir)
+        / Path("sdxl_scheduler_vmfbs")
+        / Path(sdxl_scheduler_mlir.path.name).with_suffix(f".cpu.vmfb"),
+    )
+
+
+###############################################################################
+# ROCM
+###############################################################################
+
+
+def test_compile_scheduler_rocm(sdxl_scheduler_mlir):
+    VmfbManager.sdxl_scheduler_rocm_vmfb = iree_compile(
+        sdxl_scheduler_mlir,
+        ROCM_COMPILE_FLAGS,
+        Path(vmfb_dir)
+        / Path("sdxl_scheduler_vmfbs")
+        / Path(sdxl_scheduler_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"),
+    )

experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py

@@ -10,7 +10,7 @@
 from conftest import VmfbManager
 from pathlib import Path
 
-rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
+rocm_chip = os.getenv("ROCM_CHIP", default="gfx942")
 vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd())
 
 ###############################################################################