Update

[ghstack-poisoned]

Author: GregoryComer

.ci/scripts/wheel/post_build_script_windows.sh

@@ -5,46 +5,71 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import os
-import subprocess
-import test_base
-from examples.models import Backend, Model
-from test_base import ModelTest
 from typing import List
 
-def map_backend_name(name: str) -> str:
-    # Map the backend name to the string used by the Windows test jobs, which use
-    # a slightly different convention. This is an artifact of us being mid-update
-    # of the model test logic.
-    # TODO(gjcomer) Clean this up when we update the model test CI.
+import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
+from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL
+from executorch.examples.models.model_factory import EagerModelFactory
+from executorch.examples.xnnpack import MODEL_NAME_TO_OPTIONS
+from executorch.examples.xnnpack.quantization.utils import quantize as quantize_xnn
+from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower
+from executorch.extension.pybindings.portable_lib import (
+    _load_for_executorch_from_buffer,
+)
+from test_base import ModelTest
+
+
+def test_model_xnnpack(model: Model, quantize: bool) -> None:
+    model_instance, example_inputs, _, _ = EagerModelFactory.create_model(
+        *MODEL_NAME_TO_MODEL[str(model)]
+    )
+
+    model_instance.eval()
+    ref_outputs = model_instance(*example_inputs)
+
+    if quantize:
+        quant_type = MODEL_NAME_TO_OPTIONS[str(model)].quantization
+        model_instance = torch.export.export_for_training(
+            model_instance, example_inputs
+        )
+        model_instance = quantize_xnn(
+            model_instance.module(), example_inputs, quant_type
+        )
+
+    lowered = to_edge_transform_and_lower(
+        torch.export.export(model_instance, example_inputs),
+        partitioner=[XnnpackPartitioner()],
+        compile_config=EdgeCompileConfig(
+            _check_ir_validity=False,
+        ),
+    ).to_executorch()
+
+    loaded_model = _load_for_executorch_from_buffer(lowered.buffer)
+    et_outputs = loaded_model([*example_inputs])
+
+    if isinstance(ref_outputs, torch.Tensor):
+        ref_outputs = (ref_outputs,)
+
+    assert len(ref_outputs) == len(et_outputs)
+    for i in range(len(ref_outputs)):
+        assert torch.allclose(ref_outputs[i], et_outputs[i], atol=1e-5)
 
-    if name == "xnnpack-quantization-delegation":
-        return "xnnpack-q8"
-    else:
-        return name
 
 def run_tests(model_tests: List[ModelTest]) -> None:
     for model_test in model_tests:
-        subprocess.run(
-            [
-                "powershell.exe",
-                os.path.join(test_base._repository_root_dir(), ".ci/scripts/test_model.ps1"),
-                "-ModelName",
-                str(model_test.model),
-                "-Backend",
-                map_backend_name(str(model_test.backend)),
-            ],
-            check=True,
-            cwd=test_base._repository_root_dir(),
-        )
+        if model_test.backend == Backend.Xnnpack:
+            test_model_xnnpack(model_test.model, quantize=False)
+        else:
+            raise RuntimeError(f"Unsupported backend {model_test.backend}.")
 
 
 if __name__ == "__main__":
     run_tests(
         model_tests=[
             ModelTest(
                 model=Model.Mv3,
-                backend=Backend.XnnpackQuantizationDelegation,
+                backend=Backend.Xnnpack,
             ),
         ]
     )

.ci/scripts/wheel/test_windows.py

@@ -38,7 +38,7 @@ jobs:
           - repository: pytorch/executorch
             pre-script: .ci\\scripts\\wheel\\pre_build_script.sh
             env-script: .ci\\scripts\\wheel\\vc_env_helper.bat
-            post-script: .ci\\scripts\\wheel\\post_build_script_windows.sh
+            post-script: .ci\\scripts\\wheel\\post_build_script.sh
             smoke-test-script: .ci/scripts/wheel/test_windows.py
             package-name: executorch
     name: ${{ matrix.repository }}

.github/workflows/build-wheels-windows.yml

@@ -44,6 +44,7 @@ def __str__(self) -> str:
 
 
 class Backend(str, Enum):
+    Xnnpack = ("xnnpack",)
     XnnpackQuantizationDelegation = "xnnpack-quantization-delegation"
     CoreMlExportOnly = "coreml"
     CoreMlExportAndTest = "coreml-test"  # AOT export + test with runner