test out all changes for regression testing, no benchmarking changes yet

Signed-off-by: saienduri <saimanas.enduri@amd.com>

.github/workflows/pkgci_regression_test.yml

@@ -221,10 +221,8 @@ jobs:
       - name: Install external TestSuite Python requirements
         run: |
           source ${VENV_DIR}/bin/activate
-          pip uninstall -y iree-regression-suite
           python -m pip install -r SHARK-TestSuite/iree_tests/requirements.txt
           python -m pip install -r SHARK-TestSuite/iree_special_models/requirements.txt
-          pip install --no-compile --pre --upgrade -e SHARK-TestSuite/common_tools
       - name: Download remote files for real weight model tests
         run: |
           source ${VENV_DIR}/bin/activate

experimental/regression_suite/ireers/artifacts.py

@@ -10,6 +10,7 @@
 from tqdm import tqdm
 import urllib.parse
 import urllib.request
+import os
 
 
 def show_progress(t):
@@ -26,8 +27,8 @@ def update_to(b=1, bsize=1, tsize=None):
 
 @functools.cache
 def get_artifact_root_dir() -> Path:
-    # TODO: Make configurable.
-    return Path.cwd() / "artifacts"
+    root_path = os.getenv("IREE_TEST_FILES", default=str(Path.cwd()) + "/artifacts")
+    return Path(os.path.expanduser(root_path)).resolve()
 
 
 class ArtifactGroup:
@@ -126,19 +127,97 @@ def __init__(self, group: Union[ArtifactGroup, str], url: str):
         name = Path(urllib.parse.urlparse(url).path).name
         super().__init__(group, name, FetchedArtifact._callback)
         self.url = url
+
+    def get_azure_md5(remote_file: str, azure_blob_properties: BlobProperties):
+        """Gets the content_md5 hash for a blob on Azure, if available."""
+        content_settings = azure_blob_properties.get("content_settings")
+        if not content_settings:
+            return None
+        azure_md5 = content_settings.get("content_md5")
+        if not azure_md5:
+            logger.warning(
+                f"  Remote file '{remote_file}' on Azure is missing the "
+                "'content_md5' property, can't check if local matches remote"
+            )
+        return azure_md5
+
+
+    def get_local_md5(local_file_path: Path):
+        """Gets the content_md5 hash for a lolca file, if it exists."""
+        if not local_file_path.exists() or local_file_path.stat().st_size == 0:
+            return None
+
+        with open(local_file_path) as file, mmap.mmap(
+            file.fileno(), 0, access=mmap.ACCESS_READ
+        ) as file:
+            return hashlib.md5(file).digest()
+
+    def check_azure_hashes(self: "FetchedArtifact"):
+        """
+        Checks the hashes between the local file and azure file.
+        """
+        remote_file_name = self.url.rsplit("/", 1)[-1]
+
+        # Extract path components from Azure URL to use with the Azure Storage Blobs
+        # client library for Python (https://pypi.org/project/azure-storage-blob/).
+        #
+        # For example:
+        #   https://sharkpublic.blob.core.windows.net/sharkpublic/path/to/blob.txt
+        #                                            ^           ^
+        #   account_url:    https://sharkpublic.blob.core.windows.net
+        #   container_name: sharkpublic
+        #   blob_name:      path/to/blob.txt
+        result = re.search(r"(https.+\.net)/([^/]+)/(.+)", remote_file)
+        account_url = result.groups()[0]
+        container_name = result.groups()[1]
+        blob_name = result.groups()[2]
+
+        with BlobClient(
+            account_url,
+            container_name,
+            blob_name,
+            max_chunk_get_size=1024 * 1024 * 32,  # 32 MiB
+            max_single_get_size=1024 * 1024 * 32,  # 32 MiB
+        ) as blob_client:
+            blob_properties = blob_client.get_blob_properties()
+            blob_size_str = human_readable_size(blob_properties.size)
+            azure_md5 = get_azure_md5(remote_file, blob_properties)
+
+            local_md5 = get_local_md5(self.path)
+
+            if azure_md5 and azure_md5 == local_md5:
+                print(
+                    f"  Skipping '{remote_file_name}' download ({blob_size_str}) "
+                    "- local MD5 hash matches"
+                )
+                return True
+
+            if not local_md5:
+                print(
+                    f"  Downloading '{remote_file_name}' ({blob_size_str}) "
+                    f"to '{relative_dir}'"
+                )
+                return False
+            else:
+                print(
+                    f"  Downloading '{remote_file_name}' ({blob_size_str}) "
+                    f"to '{relative_dir}' (local MD5 does not match)"
+                )
+                return False
+
 
     @staticmethod
     def _callback(self: "FetchedArtifact"):
-        print(f"Downloading {self.url} -> {self.path}", flush=True, end="")
-        with tqdm(
-            unit="B",
-            unit_scale=True,
-            unit_divisor=1024,
-            miniters=1,
-            desc=str(self.path),
-        ) as t:
-            urllib.request.urlretrieve(self.url, self.path, reporthook=show_progress(t))
-        print(f": Retrieved {self.path.stat().st_size} bytes")
+        if self.check_azure_hashes():
+            with tqdm(
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                miniters=1,
+                desc=str(self.path),
+            ) as t:
+                urllib.request.urlretrieve(self.url, self.path, reporthook=show_progress(t))
+            print(f": Retrieved {self.path.stat().st_size} bytes")
 
 
 class StreamArtifact(Artifact):

experimental/regression_suite/setup.py

@@ -20,6 +20,7 @@
         "pytest-xdist",
         "PyYAML",
         "tqdm",
+        "azure-storage-blob",
     ],
     extras_require={},
 )

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

@@ -8,14 +8,62 @@
 from ireers import *
 import os
 
-repo_root = os.getenv("TEST_SUITE_REPO_ROOT")
-current_dir = repo_root + "/iree_special_models/sd3/prompt-encoder"
 rocm_chip = os.getenv("ROCM_CHIP", default="gfx90a")
 
 ###############################################################################
 # Fixtures
 ###############################################################################
 
+sd3_clip_inference_input_0 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.0.bin",
+    group="sd3_clip_inference_input_0",
+)
+
+sd3_clip_inference_input_1 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.1.bin",
+    group="sd3_clip_inference_input_1",
+)
+
+sd3_clip_inference_input_2 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.2.bin",
+    group="sd3_clip_inference_input_2",
+)
+
+sd3_clip_inference_input_3 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.3.bin",
+    group="sd3_clip_inference_input_3",
+)
+
+sd3_clip_inference_input_4 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.4.bin",
+    group="sd3_clip_inference_input_4",
+)
+
+sd3_clip_inference_input_5 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_input.5.bin",
+    group="sd3_clip_inference_input_5",
+)
+
+sd3_clip_inference_output_0 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_output.0.bin",
+    group="sd3_clip_inference_output_0",
+)
+
+sd3_clip_inference_output_1 = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/inference_output.1.bin",
+    group="sd3_clip_inference_output_1",
+)
+
+sd3_clip_real_weights = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/real_weights.irpa",
+    group="sd3_clip_real_weights",
+)
+
+sd3_clip_mlir = fetch_source_fixture(
+    "https://sharkpublic.blob.core.windows.net/sharkpublic/sai/sd3-prompt-encoder/model.mlirbc",
+    group="sd3_clip_mlir",
+)
+
 CPU_COMPILE_FLAGS = [
     "--iree-hal-target-backends=llvm-cpu",
     "--iree-llvmcpu-target-cpu-features=host",
@@ -27,14 +75,14 @@
 ]
 
 COMMON_RUN_FLAGS = [
-    "--input=1x77x2xi64=@inference_input.0.bin",
-    "--input=1x77x2xi64=@inference_input.1.bin",
-    "--input=1x77x2xi64=@inference_input.2.bin",
-    "--input=1x77x2xi64=@inference_input.3.bin",
-    "--input=1x77x2xi64=@inference_input.4.bin",
-    "--input=1x77x2xi64=@inference_input.5.bin",
-    "--expected_output=2x154x4096xf32=@inference_output.0.bin",
-    "--expected_output=2x2048xf32=@inference_output.1.bin",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_0.path}",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_1.path}",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_2.path}",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_3.path}",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_4.path}",
+    f"--input=1x77x2xi64=@{sd3_clip_inference_input_5.path}",
+    f"--expected_output=2x154x4096xf32=@{sd3_clip_inference_output_0.path}",
+    f"--expected_output=2x2048xf32=@{sd3_clip_inference_output_1.path}",
 ]
 
 ROCM_COMPILE_FLAGS = [
@@ -54,55 +102,50 @@
     "--iree-preprocessing-pass-pipeline=builtin.module(iree-preprocessing-transpose-convolution-pipeline, util.func(iree-preprocessing-pad-to-intrinsics{pad-target-type=conv}))",
 ]
 
-mlir_path = current_dir + "/model.mlirbc"
-compile_cpu_cmd = get_compile_cmd(mlir_path, "model_cpu.vmfb", CPU_COMPILE_FLAGS)
-compile_rocm_cmd = get_compile_cmd(mlir_path, "model_rocm.vmfb", ROCM_COMPILE_FLAGS)
-
 ###############################################################################
 # CPU
 ###############################################################################
 
+cpu_vmfb = None
 
 def test_compile_clip_cpu():
-    iree_compile(mlir_path, "model_cpu.vmfb", CPU_COMPILE_FLAGS, current_dir)
+    cpu_vmfb = iree_compile(sd3_clip_mlir, "cpu", CPU_COMPILE_FLAGS)
 
 
 @pytest.mark.depends(on=["test_compile_clip_cpu"])
 def test_run_clip_cpu():
-    vmfb_path = current_dir + "/model_cpu.vmfb"
     iree_run_module(
-        vmfb_path,
-        [
-            "--device=local-task",
-            "--parameters=model=real_weights.irpa",
+        cpu_vmfb,
+        device="local-task",
+        function="encode_tokens",
+        args = [
+            f"--parameters=model={sd3_clip_real_weights.path}",
             "--expected_f32_threshold=0.15f",
         ]
-        + COMMON_RUN_FLAGS,
-        current_dir,
-        compile_cpu_cmd,
+        + COMMON_RUN_FLAGS
     )
 
 
 ###############################################################################
 # ROCM
 ###############################################################################
 
+rocm_vmfb = None
 
 @pytest.mark.xfail(
     raises=IreeCompileException,
     strict=True,
     reason="Expected compilation to fail (remove xfail for test_compile_clip_rocm)",
 )
 def test_compile_clip_rocm():
-    iree_compile(mlir_path, "model_rocm.vmfb", ROCM_COMPILE_FLAGS, current_dir)
+    rocm_vmfb = iree_compile(sd3_clip_mlir, f"rocm_{rocm_chip}", ROCM_COMPILE_FLAGS)
 
 
 @pytest.mark.depends(on=["test_compile_clip_rocm"])
 def test_run_clip_rocm():
-    vmfb_path = current_dir + "/model_rocm.vmfb"
     return iree_run_module(
-        vmfb_path,
-        ["--device=hip", "--parameters=model=real_weights.irpa"] + COMMON_RUN_FLAGS,
-        current_dir,
-        compile_rocm_cmd,
+        rocm_vmfb,
+        device="hip",
+        function="encode_tokens",
+        args=[f"--parameters=model={sd3_clip_real_weights.path}"] + COMMON_RUN_FLAGS,
     )