diff --git a/experimental/regression_suite/shark-test-suite-models/conftest.py b/experimental/regression_suite/shark-test-suite-models/conftest.py index 03f925804334..8e62bcb97274 100644 --- a/experimental/regression_suite/shark-test-suite-models/conftest.py +++ b/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 diff --git a/experimental/regression_suite/shark-test-suite-models/sd3/test_clip.py b/experimental/regression_suite/shark-test-suite-models/sd3/test_clip.py index 61368765bf6a..d5e2778481c6 100644 --- a/experimental/regression_suite/shark-test-suite-models/sd3/test_clip.py +++ b/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()) ############################################################################### diff --git a/experimental/regression_suite/shark-test-suite-models/sd3/test_mmdit.py b/experimental/regression_suite/shark-test-suite-models/sd3/test_mmdit.py index 12c25b53966d..0cbde07ba31a 100644 --- a/experimental/regression_suite/shark-test-suite-models/sd3/test_mmdit.py +++ b/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()) ############################################################################### diff --git a/experimental/regression_suite/shark-test-suite-models/sd3/test_vae.py b/experimental/regression_suite/shark-test-suite-models/sd3/test_vae.py index b41566c6ae0f..72ae9e28167e 100644 --- a/experimental/regression_suite/shark-test-suite-models/sd3/test_vae.py +++ b/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()) ############################################################################### diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py index 416f57bc3955..978595cd4c7d 100644 --- a/experimental/regression_suite/shark-test-suite-models/sdxl/test_clip.py +++ b/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()) ############################################################################### diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_scheduler.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_scheduler.py new file mode 100644 index 000000000000..2003d88db07c --- /dev/null +++ b/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"), + ) diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py index 42eb58f5d934..6eb8d903759b 100644 --- a/experimental/regression_suite/shark-test-suite-models/sdxl/test_vae.py +++ b/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()) ###############################################################################