New Feature: Support User-provided Mix Spatial Mapping, And spatial_mapping_hint Definition Can Be Provided in the Mapping File for Mix Spatial Mapping Auto-generation.

tests/main/test_with_mix_spatial_mapping/test_ascend_like.py

@@ -0,0 +1,63 @@
+import pytest
+
+from zigzag.api import get_hardware_performance_zigzag_with_mix_spatial_mapping
+
+# Test case for when spatial_mapping is provided, while spatial_mapping_hint is not provided.
+
+workloads = (
+    "zigzag/inputs/examples/workload/alexnet.onnx",
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx",
+    "zigzag/inputs/examples/workload/resnet18.onnx",
+    "zigzag.inputs.examples.workload.resnet18",
+)
+
+# Expected energy and latency for each workload defined above
+ens_lats = {
+    "zigzag/inputs/examples/workload/alexnet.onnx": (5667407342.66, 8528846),
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx": (921552096.0700004, 3828967),
+    "zigzag/inputs/examples/workload/resnet18.onnx": (1679218425.5100002, 3713386),
+    "zigzag.inputs.examples.workload.resnet18": (2290766279.31, 4442443),
+}
+
+
+@pytest.fixture
+def mapping():
+    ascend_like_mapping = {
+        "default": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("K", 16),
+                "D2": (("C", 4), ("FX", 3)),
+                "D3": ("OX", 2),
+                "D4": ("OY", 2),
+            },
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+        "Add": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("G", 16),
+                "D2": ("C", 1),
+                "D3": ("OX", 1),
+                "D4": ("OY", 1),
+            },
+            "memory_operand_links": {"O": "O", "X": "I2", "Y": "I1"},
+        },
+    }
+
+    return ascend_like_mapping
+
+
+@pytest.fixture
+def accelerator():
+    return "zigzag.inputs.examples.hardware.Ascend_like"
+
+
+@pytest.mark.parametrize("workload", workloads)
+def test_api(workload, accelerator, mapping):
+    (energy, latency, cmes) = get_hardware_performance_zigzag_with_mix_spatial_mapping(
+        workload, accelerator, mapping
+    )
+    (expected_energy, expected_latency) = ens_lats[workload]
+    assert energy == pytest.approx(expected_energy)
+    assert latency == pytest.approx(expected_latency)

tests/main/test_with_mix_spatial_mapping/test_edge_tpu_like.py

@@ -0,0 +1,79 @@
+import pytest
+
+from zigzag.api import get_hardware_performance_zigzag_with_mix_spatial_mapping
+
+# Test case for when both spatial_mapping and spatial_mapping_hint are provided.
+
+workloads = (
+    "zigzag/inputs/examples/workload/alexnet.onnx",
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx",
+    "zigzag/inputs/examples/workload/resnet18.onnx",
+    "zigzag.inputs.examples.workload.resnet18",
+)
+
+# Expected energy and latency for each workload defined above
+ens_lats = {
+    "zigzag/inputs/examples/workload/alexnet.onnx": (5582430184.085, 8343378),
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx": (762066732.5049998, 3003074),
+    "zigzag/inputs/examples/workload/resnet18.onnx": (1743190534.155, 5305825),
+    "zigzag.inputs.examples.workload.resnet18": (2087322696.315, 6155355),
+}
+
+
+@pytest.fixture
+def mapping():
+    edge_tpu_like_mapping = {
+        "default": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("K", 8),
+                "D2": (("C", 2), ("FX", 2), ("FY", 2)),
+                "D3": (("OX", 2), ("OY", 2)),
+                "D4": (("OX", 2), ("OY", 2)),
+            },
+            # spatial_mapping_hint will not work if the mapping on every dimension is provided in spatial_mapping
+            "spatial_mapping_hint": {
+                "D1": ["K"],
+                "D2": ["C", "FX", "FY"],
+                "D3": ["OX", "OY"],
+                "D4": ["OX", "OY"],
+            },
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+        "Add": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("G", 8),
+                "D2": ("C", 1),
+                "D3": ("OX", 1),
+                "D4": ("OY", 1),
+            },
+            "memory_operand_links": {"O": "O", "X": "I2", "Y": "I1"},
+        },
+        "Pooling": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("G", 8),
+                "D2": ("C", 1),
+                "D3": ("OX", 1),
+                "D4": ("OY", 1),
+            },
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+    }
+    return edge_tpu_like_mapping
+
+
+@pytest.fixture
+def accelerator():
+    return "zigzag.inputs.examples.hardware.Edge_TPU_like"
+
+
+@pytest.mark.parametrize("workload", workloads)
+def test_api(workload, accelerator, mapping):
+    (energy, latency, cmes) = get_hardware_performance_zigzag_with_mix_spatial_mapping(
+        workload, accelerator, mapping
+    )
+    (expected_energy, expected_latency) = ens_lats[workload]
+    assert energy == pytest.approx(expected_energy)
+    assert latency == pytest.approx(expected_latency)

tests/main/test_with_mix_spatial_mapping/test_meta_prototype_like.py

@@ -0,0 +1,63 @@
+import pytest
+
+from zigzag.api import get_hardware_performance_zigzag_with_mix_spatial_mapping
+
+# Test case for when an incomplete spatial_mapping is provided and spatial_mapping_hint is also provided.
+
+workloads = (
+    "zigzag/inputs/examples/workload/alexnet.onnx",
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx",
+    "zigzag/inputs/examples/workload/resnet18.onnx",
+    "zigzag.inputs.examples.workload.resnet18",
+)
+
+# Expected energy and latency for each workload defined above
+ens_lats = {
+    "zigzag/inputs/examples/workload/alexnet.onnx": (5681909351.240001, 8299150),
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx": (919452681.2249999, 2894129),
+    "zigzag/inputs/examples/workload/resnet18.onnx": (1789888904.4450002, 3472280),
+    "zigzag.inputs.examples.workload.resnet18": (2348207081.7949996, 4238517),
+}
+
+
+@pytest.fixture
+def mapping():
+    meta_prototype_like_mapping = {
+        "default": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("K", 32),
+                # "D2": ("C", 2),
+                "D3": (("OX", 2), ("OY", 2)),
+                "D4": (("OX", 2), ("OY", 2)),
+            },
+            "spatial_mapping_hint": {"D2": ["C"]},
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+        "Add": {
+            "core_allocation": 1,
+            "spatial_mapping": {
+                "D1": ("G", 32),
+                "D2": ("C", 1),
+                "D3": ("OX", 1),
+                "D4": ("OY", 1),
+            },
+            "memory_operand_links": {"O": "O", "X": "I2", "Y": "I1"},
+        },
+    }
+    return meta_prototype_like_mapping
+
+
+@pytest.fixture
+def accelerator():
+    return "zigzag.inputs.examples.hardware.Meta_prototype"
+
+
+@pytest.mark.parametrize("workload", workloads)
+def test_api(workload, accelerator, mapping):
+    (energy, latency, cmes) = get_hardware_performance_zigzag_with_mix_spatial_mapping(
+        workload, accelerator, mapping
+    )
+    (expected_energy, expected_latency) = ens_lats[workload]
+    assert energy == pytest.approx(expected_energy)
+    assert latency == pytest.approx(expected_latency)

tests/main/test_with_mix_spatial_mapping/test_tesla_npu_like.py

@@ -0,0 +1,59 @@
+import pytest
+
+from zigzag.api import get_hardware_performance_zigzag_with_mix_spatial_mapping
+
+# Test case for when only spatial_mapping_hint is provided.
+
+workloads = (
+    "zigzag/inputs/examples/workload/alexnet.onnx",
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx",
+    "zigzag/inputs/examples/workload/resnet18.onnx",
+    "zigzag.inputs.examples.workload.resnet18",
+)
+
+# Expected energy and latency for each workload defined above
+ens_lats = {
+    "zigzag/inputs/examples/workload/alexnet.onnx": (6040086796.366001, 8389669),
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx": (958401881.3470002, 1964453),
+    "zigzag/inputs/examples/workload/resnet18.onnx": (1724869681.4799998, 3257898),
+    "zigzag.inputs.examples.workload.resnet18": (2220861655.6660004, 3934616),
+}
+
+
+@pytest.fixture
+def mapping():
+    tesla_npu_like_mapping = {
+        "default": {
+            "core_allocation": 1,
+            # "spatial_mapping": {"D1": ("K", 32), "D2": ("OX", 8), "D3": ("OY", 4)},
+            "spatial_mapping_hint": {"D1": ["K"], "D2": ["OX", "OY"]},
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+        "Add": {
+            "core_allocation": 1,
+            "spatial_mapping": {"D1": ("G", 32), "D2": ("OX", 1), "D3": ("OY", 1)},
+            "memory_operand_links": {"O": "O", "X": "I2", "Y": "I1"},
+        },
+        "Pooling": {
+            "core_allocation": 1,
+            "spatial_mapping": {"D1": ("G", 32), "D2": ("OX", 1), "D3": ("OY", 1)},
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+    }
+
+    return tesla_npu_like_mapping
+
+
+@pytest.fixture
+def accelerator():
+    return "zigzag.inputs.examples.hardware.Tesla_NPU_like"
+
+
+@pytest.mark.parametrize("workload", workloads)
+def test_api(workload, accelerator, mapping):
+    (energy, latency, cmes) = get_hardware_performance_zigzag_with_mix_spatial_mapping(
+        workload, accelerator, mapping
+    )
+    (expected_energy, expected_latency) = ens_lats[workload]
+    assert energy == pytest.approx(expected_energy)
+    assert latency == pytest.approx(expected_latency)

tests/main/test_with_mix_spatial_mapping/test_tpu_like.py

@@ -0,0 +1,68 @@
+import pytest
+
+from zigzag.api import get_hardware_performance_zigzag_with_mix_spatial_mapping
+
+# Test case for when more non-existent dimensions are provided in spatial_mapping_hint.
+
+workloads = (
+    "zigzag/inputs/examples/workload/alexnet.onnx",
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx",
+    "zigzag/inputs/examples/workload/resnet18.onnx",
+    "zigzag.inputs.examples.workload.resnet18",
+)
+
+# Expected energy and latency for each workload defined above
+ens_lats = {
+    "zigzag/inputs/examples/workload/alexnet.onnx": (5468347515.336, 8210374),
+    "zigzag/inputs/examples/workload/mobilenetv2.onnx": (730691567.0230001, 3358406),
+    "zigzag/inputs/examples/workload/resnet18.onnx": (1648700710.272, 2946593),
+    "zigzag.inputs.examples.workload.resnet18": (1972279074.768, 3455539),
+}
+
+
+@pytest.fixture
+def mapping():
+    tpu_like_mapping = {
+        "default": {
+            "core_allocation": 1,
+            # "spatial_mapping": {
+            #     "D1": ("K", 32),
+            #     "D2": (("C", 2), ("FX", 3), ("FY", 3)),
+            # },
+            # D3 and D4 in spatial_mapping_hint will not work, since they do not exist in the hardware dimensions.
+            "spatial_mapping_hint": {
+                "D1": ["K"],
+                "D2": ["C", "FX", "FY"],
+                "D3": ["K", "OX"],
+                "D4": ["OX", "OY"],
+            },
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+        "Add": {
+            "core_allocation": 1,
+            "spatial_mapping": {"D1": ("G", 32), "D2": ("C", 1)},
+            "memory_operand_links": {"O": "O", "X": "I2", "Y": "I1"},
+        },
+        "Pooling": {
+            "core_allocation": 1,
+            "spatial_mapping": {"D1": ("G", 32), "D2": ("C", 1)},
+            "memory_operand_links": {"O": "O", "W": "I2", "I": "I1"},
+        },
+    }
+
+    return tpu_like_mapping
+
+
+@pytest.fixture
+def accelerator():
+    return "zigzag.inputs.examples.hardware.TPU_like"
+
+
+@pytest.mark.parametrize("workload", workloads)
+def test_api(workload, accelerator, mapping):
+    (energy, latency, cmes) = get_hardware_performance_zigzag_with_mix_spatial_mapping(
+        workload, accelerator, mapping
+    )
+    (expected_energy, expected_latency) = ens_lats[workload]
+    assert energy == pytest.approx(expected_energy)
+    assert latency == pytest.approx(expected_latency)