From cc44a855465ec783880275ff2168cd753c40b384 Mon Sep 17 00:00:00 2001 From: saienduri <77521230+saienduri@users.noreply.github.com> Date: Fri, 23 Aug 2024 17:20:18 -0700 Subject: [PATCH] Rework special model testing to avoid shared cache interference. (#18344) This commit reworks the special model testing for sdxl and sd3 so that different instances generating compiler artifacts don't write to the same shared cache. This is a major problem because the things we are running for each test are no longer consistent. On top of that, benchmarking was picking up from this same cache, so there were a bunch of failures leading to chaos. Now, the output artifacts are written to a place local to each instance running, so there should be no attempts to write to the same cache in between jobs. We still could use a better longterm solution to avoid this problem for our input artifacts (mlir, weights), but for now we are good as they rarely change and we can be strategic with when we update them. Huggingface seems like a good option to move toward Fixes https://github.com/iree-org/iree/issues/18336 --------- Signed-off-by: saienduri --- .github/workflows/pkgci_regression_test.yml | 1 + .../benchmarks/sdxl/benchmark_sdxl_rocm.py | 28 ++++---- .../regression_suite/ireers_tools/fixtures.py | 67 +++++++++---------- .../shark-test-suite-models/sd3/test_clip.py | 14 +++- .../shark-test-suite-models/sd3/test_mmdit.py | 13 +++- .../shark-test-suite-models/sd3/test_vae.py | 16 ++++- .../shark-test-suite-models/sdxl/test_clip.py | 14 +++- .../shark-test-suite-models/sdxl/test_unet.py | 27 ++++++-- .../shark-test-suite-models/sdxl/test_vae.py | 14 +++- 9 files changed, 129 insertions(+), 65 deletions(-) diff --git a/.github/workflows/pkgci_regression_test.yml b/.github/workflows/pkgci_regression_test.yml index 588f3bfbacfd..22ce023753ed 100644 --- a/.github/workflows/pkgci_regression_test.yml +++ b/.github/workflows/pkgci_regression_test.yml @@ -251,6 +251,7 @@ jobs: PACKAGE_DOWNLOAD_DIR: ${{ github.workspace }}/.packages IREE_TEST_PATH_EXTENSION: ${{ github.workspace }}/build_tools/pkgci/external_test_suite VENV_DIR: ${{ github.workspace }}/venv + TEST_OUTPUT_ARTIFACTS: ${{ github.workspace }}/model_output_artifacts steps: - name: Checking out IREE repository uses: actions/checkout@v4.1.7 diff --git a/experimental/benchmarks/sdxl/benchmark_sdxl_rocm.py b/experimental/benchmarks/sdxl/benchmark_sdxl_rocm.py index c674e5ae47bb..3c7c2dd2bef3 100644 --- a/experimental/benchmarks/sdxl/benchmark_sdxl_rocm.py +++ b/experimental/benchmarks/sdxl/benchmark_sdxl_rocm.py @@ -14,12 +14,16 @@ import tabulate from pytest_check import check +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) benchmark_dir = os.path.dirname(os.path.realpath(__file__)) artifacts_dir = os.getenv("IREE_TEST_FILES", default=Path.cwd()) + "/artifacts" artifacts_dir = Path(os.path.expanduser(artifacts_dir)).resolve() prompt_encoder_dir = f"{artifacts_dir}/sdxl_clip" scheduled_unet_dir = f"{artifacts_dir}/sdxl_unet" vae_decode_dir = f"{artifacts_dir}/sdxl_vae" +prompt_encoder_dir_compile = f"{vmfb_dir}/sdxl_clip_vmfbs" +scheduled_unet_dir_compile = f"{vmfb_dir}/sdxl_unet_vmfbs" +vae_decode_dir_compile = f"{vmfb_dir}/sdxl_vae_vmfbs" def run_iree_command(args: Sequence[str] = ()): @@ -69,11 +73,11 @@ def run_sdxl_rocm_benchmark(rocm_chip): "iree-benchmark-module", f"--device=hip", "--device_allocator=caching", - f"--module={prompt_encoder_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={prompt_encoder_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={prompt_encoder_dir}/real_weights.irpa", - f"--module={scheduled_unet_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={scheduled_unet_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={scheduled_unet_dir}/real_weights.irpa", - f"--module={vae_decode_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={vae_decode_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={vae_decode_dir}/real_weights.irpa", f"--module={benchmark_dir}/sdxl_full_pipeline_fp16_rocm.vmfb", "--function=tokens_to_image", @@ -95,7 +99,7 @@ def run_sdxl_unet_rocm_benchmark(rocm_chip): "iree-benchmark-module", f"--device=hip", "--device_allocator=caching", - f"--module={scheduled_unet_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={scheduled_unet_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={scheduled_unet_dir}/real_weights.irpa", "--function=run_forward", "--input=1x4x128x128xf16", @@ -116,7 +120,7 @@ def run_sdxl_prompt_encoder_rocm_benchmark(rocm_chip): "iree-benchmark-module", f"--device=hip", "--device_allocator=caching", - f"--module={prompt_encoder_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={prompt_encoder_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={prompt_encoder_dir}/real_weights.irpa", "--function=encode_prompts", "--input=1x64xi64", @@ -135,7 +139,7 @@ def run_sdxl_vae_decode_rocm_benchmark(rocm_chip): "iree-benchmark-module", f"--device=hip", "--device_allocator=caching", - f"--module={vae_decode_dir}/model.rocm_{rocm_chip}.vmfb", + f"--module={vae_decode_dir_compile}/model.rocm_{rocm_chip}.vmfb", f"--parameters=model={vae_decode_dir}/real_weights.irpa", "--function=main", "--input=1x4x128x128xf16", @@ -221,7 +225,7 @@ def test_sdxl_rocm_benchmark( logging.getLogger().info(mean_line) # unet compilation stats check - with open(f"{scheduled_unet_dir}/compilation_info.json", "r") as file: + with open(f"{scheduled_unet_dir_compile}/compilation_info.json", "r") as file: comp_stats = json.load(file) unet_dispatch_count = int( comp_stats["stream-aggregate"]["execution"]["dispatch-count"] @@ -232,7 +236,7 @@ def test_sdxl_rocm_benchmark( ) logging.getLogger().info(compilation_line) - module_path = f"{scheduled_unet_dir}/model.rocm_{rocm_chip}.vmfb" + module_path = f"{scheduled_unet_dir_compile}/model.rocm_{rocm_chip}.vmfb" unet_binary_size = Path(module_path).stat().st_size compilation_line = ( f"Scheduled Unet Binary Size: {unet_binary_size} bytes" @@ -250,7 +254,7 @@ def test_sdxl_rocm_benchmark( logging.getLogger().info(mean_line) # prompt encoder compilation stats check - with open(f"{prompt_encoder_dir}/compilation_info.json", "r") as file: + with open(f"{prompt_encoder_dir_compile}/compilation_info.json", "r") as file: comp_stats = json.load(file) clip_dispatch_count = int( comp_stats["stream-aggregate"]["execution"]["dispatch-count"] @@ -261,7 +265,7 @@ def test_sdxl_rocm_benchmark( ) logging.getLogger().info(compilation_line) - module_path = f"{prompt_encoder_dir}/model.rocm_{rocm_chip}.vmfb" + module_path = f"{prompt_encoder_dir_compile}/model.rocm_{rocm_chip}.vmfb" clip_binary_size = Path(module_path).stat().st_size compilation_line = ( f"Prompt Encoder Binary Size: {clip_binary_size} bytes" @@ -279,7 +283,7 @@ def test_sdxl_rocm_benchmark( logging.getLogger().info(mean_line) # vae decode compilation stats check - with open(f"{vae_decode_dir}/compilation_info.json", "r") as file: + with open(f"{vae_decode_dir_compile}/compilation_info.json", "r") as file: comp_stats = json.load(file) vae_dispatch_count = int( comp_stats["stream-aggregate"]["execution"]["dispatch-count"] @@ -290,7 +294,7 @@ def test_sdxl_rocm_benchmark( ) logging.getLogger().info(compilation_line) - module_path = f"{vae_decode_dir}/model.rocm_{rocm_chip}.vmfb" + module_path = f"{vae_decode_dir_compile}/model.rocm_{rocm_chip}.vmfb" vae_binary_size = Path(module_path).stat().st_size compilation_line = ( f"VAE Decode Binary Size: {vae_binary_size} bytes" diff --git a/experimental/regression_suite/ireers_tools/fixtures.py b/experimental/regression_suite/ireers_tools/fixtures.py index 543f64ecf011..6d134c770323 100644 --- a/experimental/regression_suite/ireers_tools/fixtures.py +++ b/experimental/regression_suite/ireers_tools/fixtures.py @@ -10,6 +10,7 @@ from pathlib import Path import subprocess import time +import os from .artifacts import ( Artifact, @@ -35,60 +36,52 @@ def fetcher() -> FetchedArtifact: return fetcher -def iree_compile(source: Artifact, compiled_variant: str, flags: Sequence[str]): - name = Path(source.name).with_suffix(f".{compiled_variant}.vmfb") - - def callback(vmfb_artifact: ProducedArtifact): - sep = "\n " - print("**************************************************************") - print(f"Compiling {source} -> {vmfb_artifact} with flags:") - print(f" {sep.join(flags)}") - exec_args = ( - [ - "iree-compile", - "-o", - str(vmfb_artifact.path), - str(source.path), - ] - + IREE_COMPILE_QOL_FLAGS - + flags - ) - start_time = time.time() - subprocess.run( - exec_args, check=True, capture_output=True, cwd=source.group.directory - ) - run_time = time.time() - start_time - print(f"Compilation succeeded in {run_time}s") - print("**************************************************************") - - return ProducedArtifact(source.group, name, callback, depends=[source]).start() +def iree_compile(source: Artifact, flags: Sequence[str], vmfb_path: Path): + if not os.path.exists(vmfb_path.parent): + os.makedirs(vmfb_path.parent) + sep = "\n " + print("**************************************************************") + print(f" {sep.join(flags)}") + exec_args = ( + [ + "iree-compile", + "-o", + str(vmfb_path), + str(source.path), + ] + + IREE_COMPILE_QOL_FLAGS + + flags + ) + print("Exec:", " ".join(exec_args)) + start_time = time.time() + subprocess.run(exec_args, check=True, capture_output=True, cwd=vmfb_path.parent) + run_time = time.time() - start_time + print(f"Compilation succeeded in {run_time}s") + print("**************************************************************") + return vmfb_path -def iree_run_module(vmfb: Artifact, *, device, function, args: Sequence[str] = ()): - vmfb.join() +def iree_run_module(vmfb: Path, *, device, function, args: Sequence[str] = ()): exec_args = [ "iree-run-module", f"--device={device}", - f"--module={vmfb.path}", + f"--module={vmfb}", f"--function={function}", ] exec_args.extend(args) print("**************************************************************") print("Exec:", " ".join(exec_args)) - subprocess.run(exec_args, check=True, capture_output=True, cwd=vmfb.group.directory) + subprocess.run(exec_args, check=True, capture_output=True, cwd=vmfb.parent) -def iree_benchmark_module( - vmfb: Artifact, *, device, function, args: Sequence[str] = () -): - vmfb.join() +def iree_benchmark_module(vmfb: Path, *, device, function, args: Sequence[str] = ()): exec_args = [ "iree-benchmark-module", f"--device={device}", - f"--module={vmfb.path}", + f"--module={vmfb}", f"--function={function}", ] exec_args.extend(args) print("**************************************************************") print("Exec:", " ".join(exec_args)) - subprocess.check_call(exec_args, cwd=vmfb.group.directory) + subprocess.check_call(exec_args, cwd=vmfb.parent) 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 1ba461c207ca..997dc3f710fd 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 @@ -8,8 +8,10 @@ from ireers_tools import * import os from conftest import VmfbManager +from pathlib import Path rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) ############################################################################### # Fixtures @@ -123,7 +125,11 @@ def SD3_CLIP_COMMON_RUN_FLAGS( def test_compile_clip_cpu(sd3_clip_mlir): VmfbManager.sd3_clip_cpu_vmfb = iree_compile( - sd3_clip_mlir, "cpu", CPU_COMPILE_FLAGS + sd3_clip_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_clip_vmfbs") + / Path(sd3_clip_mlir.path.name).with_suffix(f".cpu.vmfb"), ) @@ -152,7 +158,11 @@ def test_run_clip_cpu(SD3_CLIP_COMMON_RUN_FLAGS, sd3_clip_real_weights): ) def test_compile_clip_rocm(sd3_clip_mlir): VmfbManager.sd3_clip_rocm_vmfb = iree_compile( - sd3_clip_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sd3_clip_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_clip_vmfbs") + / Path(sd3_clip_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), ) 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 8e839fd4560c..5fc96c753b53 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 @@ -12,6 +12,7 @@ iree_test_path_extension = os.getenv("IREE_TEST_PATH_EXTENSION", default=Path.cwd()) rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) ############################################################################### # Fixtures @@ -107,7 +108,11 @@ def SD3_MMDIT_COMMON_RUN_FLAGS( def test_compile_mmdit_cpu(sd3_mmdit_mlir): VmfbManager.sd3_mmdit_cpu_vmfb = iree_compile( - sd3_mmdit_mlir, "cpu", CPU_COMPILE_FLAGS + sd3_mmdit_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_mmdit_vmfbs") + / Path(sd3_mmdit_mlir.path.name).with_suffix(f".cpu.vmfb"), ) @@ -137,7 +142,11 @@ def test_run_mmdit_cpu(SD3_MMDIT_COMMON_RUN_FLAGS, sd3_mmdit_real_weights): ) def test_compile_mmdit_rocm(sd3_mmdit_mlir): VmfbManager.sd3_mmdit_rocm_vmfb = iree_compile( - sd3_mmdit_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sd3_mmdit_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_mmdit_vmfbs") + / Path(sd3_mmdit_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), ) 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 5367390defe5..8037ef008f5e 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 @@ -8,8 +8,10 @@ from ireers_tools import * import os from conftest import VmfbManager +from pathlib import Path rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) ############################################################################### # Fixtures @@ -77,7 +79,13 @@ def SD3_VAE_COMMON_RUN_FLAGS( def test_compile_vae_cpu(sd3_vae_mlir): - VmfbManager.sd3_vae_cpu_vmfb = iree_compile(sd3_vae_mlir, "cpu", CPU_COMPILE_FLAGS) + VmfbManager.sd3_vae_cpu_vmfb = iree_compile( + sd3_vae_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_vae_vmfbs") + / Path(sd3_vae_mlir.path.name).with_suffix(f".cpu.vmfb"), + ) @pytest.mark.depends(on=["test_compile_vae_cpu"]) @@ -101,7 +109,11 @@ def test_run_vae_cpu(SD3_VAE_COMMON_RUN_FLAGS, sd3_vae_real_weights): def test_compile_vae_rocm(sd3_vae_mlir): VmfbManager.sd3_vae_rocm_vmfb = iree_compile( - sd3_vae_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sd3_vae_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sd3_vae_vmfbs") + / Path(sd3_vae_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), ) 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 6ed745051bbe..1eb686df245e 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 @@ -8,8 +8,10 @@ from ireers_tools import * import os from conftest import VmfbManager +from pathlib import Path rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) ############################################################################### # Fixtures @@ -111,7 +113,11 @@ def SDXL_CLIP_COMMON_RUN_FLAGS( def test_compile_clip_cpu(sdxl_clip_mlir): VmfbManager.sdxl_clip_cpu_vmfb = iree_compile( - sdxl_clip_mlir, "cpu", CPU_COMPILE_FLAGS + sdxl_clip_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_clip_vmfbs") + / Path(sdxl_clip_mlir.path.name).with_suffix(f".cpu.vmfb"), ) @@ -136,7 +142,11 @@ def test_run_clip_cpu(SDXL_CLIP_COMMON_RUN_FLAGS, sdxl_clip_real_weights): def test_compile_clip_rocm(sdxl_clip_mlir): VmfbManager.sdxl_clip_rocm_vmfb = iree_compile( - sdxl_clip_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sdxl_clip_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_clip_vmfbs") + / Path(sdxl_clip_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), ) diff --git a/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py b/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py index ef2fca97dcc1..0a5af53ae9af 100644 --- a/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py +++ b/experimental/regression_suite/shark-test-suite-models/sdxl/test_unet.py @@ -12,6 +12,7 @@ from pathlib import Path iree_test_path_extension = os.getenv("IREE_TEST_PATH_EXTENSION", default=Path.cwd()) +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") ############################################################################### @@ -123,14 +124,20 @@ def SDXL_UNET_COMMON_RUN_FLAGS( def test_compile_unet_pipeline_cpu(sdxl_unet_pipeline_mlir): VmfbManager.sdxl_unet_cpu_pipeline_vmfb = iree_compile( sdxl_unet_pipeline_mlir, - "cpu", CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_unet_vmfbs") + / Path(sdxl_unet_pipeline_mlir.path.name).with_suffix(f".cpu.vmfb"), ) def test_compile_unet_cpu(sdxl_unet_mlir): VmfbManager.sdxl_unet_cpu_vmfb = iree_compile( - sdxl_unet_mlir, "cpu", CPU_COMPILE_FLAGS + sdxl_unet_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_unet_vmfbs") + / Path(sdxl_unet_mlir.path.name).with_suffix(f".cpu.vmfb"), ) @@ -142,7 +149,7 @@ def test_run_unet_cpu(SDXL_UNET_COMMON_RUN_FLAGS, sdxl_unet_real_weights): function="produce_image_latents", args=[ f"--parameters=model={sdxl_unet_real_weights.path}", - f"--module={VmfbManager.sdxl_unet_cpu_pipeline_vmfb.path}", + f"--module={VmfbManager.sdxl_unet_cpu_pipeline_vmfb}", "--expected_f16_threshold=0.8f", ] + SDXL_UNET_COMMON_RUN_FLAGS, @@ -157,14 +164,22 @@ def test_run_unet_cpu(SDXL_UNET_COMMON_RUN_FLAGS, sdxl_unet_real_weights): def test_compile_unet_pipeline_rocm(sdxl_unet_pipeline_mlir): VmfbManager.sdxl_unet_rocm_pipeline_vmfb = iree_compile( sdxl_unet_pipeline_mlir, - f"rocm_{rocm_chip}", ROCM_PIPELINE_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_unet_vmfbs") + / Path(sdxl_unet_pipeline_mlir.path.name).with_suffix( + f".rocm_{rocm_chip}.vmfb" + ), ) def test_compile_unet_rocm(sdxl_unet_mlir): VmfbManager.sdxl_unet_rocm_vmfb = iree_compile( - sdxl_unet_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sdxl_unet_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_unet_vmfbs") + / Path(sdxl_unet_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), ) @@ -176,7 +191,7 @@ def test_run_unet_rocm(SDXL_UNET_COMMON_RUN_FLAGS, sdxl_unet_real_weights): function="produce_image_latents", args=[ f"--parameters=model={sdxl_unet_real_weights.path}", - f"--module={VmfbManager.sdxl_unet_rocm_pipeline_vmfb.path}", + f"--module={VmfbManager.sdxl_unet_rocm_pipeline_vmfb}", "--expected_f16_threshold=0.705f", ] + SDXL_UNET_COMMON_RUN_FLAGS, 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 bca80733e31e..515dd9999a08 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 @@ -8,8 +8,10 @@ from ireers_tools import * import os from conftest import VmfbManager +from pathlib import Path rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a") +vmfb_dir = os.getenv("TEST_OUTPUT_ARTIFACTS", default=Path.cwd()) ############################################################################### # Fixtures @@ -80,7 +82,11 @@ def SDXL_VAE_COMMON_RUN_FLAGS( def test_compile_vae_cpu(sdxl_vae_mlir): VmfbManager.sdxl_vae_cpu_vmfb = iree_compile( - sdxl_vae_mlir, "cpu", CPU_COMPILE_FLAGS + sdxl_vae_mlir, + CPU_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_vae_vmfbs") + / Path(sdxl_vae_mlir.path.name).with_suffix(f".cpu.vmfb"), ) @@ -105,7 +111,11 @@ def test_run_vae_cpu(SDXL_VAE_COMMON_RUN_FLAGS, sdxl_vae_real_weights): def test_compile_vae_rocm(sdxl_vae_mlir): VmfbManager.sdxl_vae_rocm_vmfb = iree_compile( - sdxl_vae_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS + sdxl_vae_mlir, + ROCM_COMPILE_FLAGS, + Path(vmfb_dir) + / Path("sdxl_vae_vmfbs") + / Path(sdxl_vae_mlir.path.name).with_suffix(f".rocm_{rocm_chip}.vmfb"), )